This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 2cfb504c12 [python] Refuse a read whose query auth rules cannot be
bound to columns (#10018)
2cfb504c12 is described below
commit 2cfb504c12feb7b0481b9bdcb9a29932f4a91c0e
Author: Jiajia Li <[email protected]>
AuthorDate: Mon Sep 21 10:40:18 2026 +0800
[python] Refuse a read whose query auth rules cannot be bound to columns
(#10018)
---
paimon-python/pypaimon/catalog/table_query_auth.py | 210 +++++
.../pypaimon/common/predicate_json_parser.py | 9 +-
.../pypaimon/read/streaming_table_scan.py | 5 +-
paimon-python/pypaimon/read/table_read.py | 27 +
paimon-python/pypaimon/read/table_scan.py | 37 +-
.../pypaimon/tests/deferred_blob_resolve_test.py | 13 +-
.../pypaimon/tests/multimodal_temporal_test.py | 47 +-
.../pypaimon/tests/query_auth_validation_test.py | 957 +++++++++++++++++++++
.../pypaimon/tests/read_builder_explain_test.py | 3 +-
9 files changed, 1267 insertions(+), 41 deletions(-)
diff --git a/paimon-python/pypaimon/catalog/table_query_auth.py
b/paimon-python/pypaimon/catalog/table_query_auth.py
index 523c42742d..07cb260e75 100644
--- a/paimon-python/pypaimon/catalog/table_query_auth.py
+++ b/paimon-python/pypaimon/catalog/table_query_auth.py
@@ -16,18 +16,74 @@
# limitations under the License.
################################################################################
+import json
from typing import Callable, Dict, List, Optional
import pyarrow as pa
import pyarrow.compute as pc
from pypaimon.common.predicate_json_parser import (
+ _collect_all_field_refs_from_transform,
extract_referenced_fields,
parse_predicate_to_batch_filter,
)
from pypaimon.read.plan import Plan
from pypaimon.read.query_auth_split import QueryAuthSplit
from pypaimon.schema.data_types import DataField
+from pypaimon.table.special_fields import SpecialFields
+
+# SchemaValidation forbids a stored column from carrying one of these names,
so such a rule
+# always means the synthetic field, which binds only where the query already
projects it.
+_RESERVED_RULE_NAMES = SpecialFields.SYSTEM_FIELD_NAMES
+
+_SYSTEM_FIELD_IDS = {
+ field.name: field.id
+ for field in (SpecialFields.SEQUENCE_NUMBER,
+ SpecialFields.VALUE_KIND,
+ SpecialFields.ROW_ID)
+}
+
+
+def _check_field_exists(
+ rule: str,
+ field: str,
+ table_field_names: List[str],
+):
+ if field in _RESERVED_RULE_NAMES:
+ return
+ if field not in table_field_names:
+ raise ValueError(
+ "{} references column '{}' which does not exist in table schema
{}. "
+ "The rule may be stale after a column rename or drop; refusing to
read."
+ .format(rule, field, table_field_names))
+
+
+def _check_not_renamed(
+ rule: str,
+ field: str,
+ latest_by_name: Dict[str, DataField],
+ read_by_name: Dict[str, DataField],
+ read_by_id: Dict[int, DataField],
+):
+ # absent from latest: already reported by validate_against_schema
+ if field not in latest_by_name:
+ return
+ field_id = latest_by_name[field].id
+ read_field = read_by_name.get(field)
+ if read_field is not None:
+ # a dropped and re-added column keeps the name but gets a fresh id
+ if read_field.id != field_id:
+ raise ValueError(
+ "{} references column '{}' which the snapshot being read
exposes as a "
+ "different column of the same name (dropped and re-added
since); refusing "
+ "to read to avoid applying the rule to unrelated
data.".format(rule, field))
+ return
+ renamed = read_by_id.get(field_id)
+ if renamed is not None:
+ raise ValueError(
+ "{} references column '{}' which the snapshot being read exposes
as '{}' "
+ "(renamed since); refusing to read to avoid applying the rule by a
stale name."
+ .format(rule, field, renamed.name))
def reject_search_under_query_auth(table) -> None:
@@ -74,6 +130,160 @@ class TableQueryAuthResult:
return result
return combined
+ def parsed_column_masking(self) -> Dict[str, dict]:
+ if not self.column_masking:
+ return {}
+ parsed = {}
+ for column, transform_json in self.column_masking.items():
+ transform = json.loads(transform_json)
+ if transform is None:
+ raise ValueError(
+ "Column masking on '{}' is JSON null; refusing to read
rather than "
+ "returning the column unmasked.".format(column))
+ parsed[column] = transform
+ return parsed
+
+ def filter_field_names(self) -> set:
+ names = set()
+ for json_str in self.filter or []:
+ names.update(extract_referenced_fields(json_str))
+ return names
+
+ def rule_field_names(self, readable: Optional[set] = None) -> set:
+ names = self.filter_field_names()
+ for target, transform in self.parsed_column_masking().items():
+ names.add(target)
+ # an unreadable target makes the mask inert, so its inputs bind
nothing
+ if readable is None or target in readable:
+ names.update(_collect_all_field_refs_from_transform(transform))
+ return names
+
+ def validate_against_schema(
+ self,
+ table_fields: List[DataField],
+ ):
+ """Validate that every column the auth rules reference exists in the
table's latest
+ schema. A rule keyed by a since-renamed column would silently stop
masking; fail closed.
+ """
+ masking = self.parsed_column_masking()
+ table_field_names = [f.name for f in table_fields]
+ for target, transform in masking.items():
+ _check_field_exists("Column masking", target, table_field_names)
+ for source in _collect_all_field_refs_from_transform(transform):
+ _check_field_exists("Column masking", source,
table_field_names)
+ # a transform reads the raw row, so a masked input would be
published
+ # unmasked through this target
+ if source != target and source in masking:
+ raise ValueError(
+ "Column masking on '{}' reads column '{}', which is
masked too. "
+ "The mask would be computed from the raw value of '{}'
and expose "
+ "it through '{}'.".format(target, source, source,
target))
+ for operand in self.filter_field_names():
+ _check_field_exists("Row filter", operand, table_field_names)
+
+ def validate_readable_without_rename(
+ self,
+ latest_fields: List[DataField],
+ snapshot_fields: List[DataField],
+ ):
+ """Fail closed when the snapshot being read exposes a rule's column
under a different
+ name, where enforcing by name would skip it. A column absent from it
stays inert.
+ """
+ latest_by_name = {f.name: f for f in latest_fields}
+ read_by_name = {f.name: f for f in snapshot_fields}
+ read_by_id = {f.id: f for f in snapshot_fields}
+ read_names = set(read_by_name)
+ for target, transform in self.parsed_column_masking().items():
+ _check_not_renamed("Column masking", target, latest_by_name,
read_by_name, read_by_id)
+ # a mask whose target is absent from the read schema is inert;
skip its inputs
+ if target in read_names:
+ for source in
_collect_all_field_refs_from_transform(transform):
+ _check_not_renamed(
+ "Column masking", source, latest_by_name,
read_by_name, read_by_id)
+ for operand in self.filter_field_names():
+ _check_not_renamed("Row filter", operand, latest_by_name,
read_by_name, read_by_id)
+
+ def validate_read_type(
+ self,
+ table_fields: List[DataField],
+ read_fields: List[DataField],
+ nested_name_paths: Optional[List[List[str]]] = None,
+ snapshot_fields: Optional[List[DataField]] = None,
+ ):
+ """Validate a physical projection before adding columns required by
authorization
+ rules: a rule must never be applied to a partially projected column.
+ """
+ read_names = {f.name for f in read_fields}
+ rule_fields = self.rule_field_names(read_names)
+ # a system column is synthesised, not stored, so it cannot be added to
the projection;
+ # a mask on an unprojected one needs nothing and stays inert
+ needed_system = {
+ name for name in self.filter_field_names() if name in
_RESERVED_RULE_NAMES}
+ for target, transform in self.parsed_column_masking().items():
+ if target not in read_names:
+ continue
+ needed_system.update(
+ name for name in
_collect_all_field_refs_from_transform(transform)
+ if name in _RESERVED_RULE_NAMES)
+ unbound = needed_system - read_names
+ if unbound:
+ raise RuntimeError(
+ "Query auth rules need system column(s) {} which the query
does not project; "
+ "a system column cannot be added to the
read.".format(sorted(unbound)))
+ # identity from the schema the rules name, type from the one the split
was written
+ # under, which may since have been widened
+ table_by_name = {f.name: f for f in table_fields}
+ type_by_name = {f.name: f for f in (snapshot_fields or table_fields)}
+ expected_id = {n: table_by_name[n].id for n in rule_fields if n in
table_by_name}
+ # a system column is absent from the table schema, so bind it by its
canonical id:
+ # a projected subfield flattens to a top-level name and could
otherwise pass as one
+ expected_id.update(
+ {n: i for n, i in _SYSTEM_FIELD_IDS.items() if n in rule_fields})
+ # a masked column reached through a generated name (email -> email__0)
is published
+ # raw; a filter operand still binds by its own name
+ masked_ids = {expected_id[n] for n in self.parsed_column_masking() if
n in expected_id}
+ for field in read_fields:
+ if field.name not in rule_fields:
+ if field.id in masked_ids:
+ raise RuntimeError(
+ "Query auth rules mask field id {}, which the query
also projects "
+ "under the generated name '{}'; the second copy would
be raw."
+ .format(field.id, field.name))
+ continue
+ want = expected_id.get(field.name)
+ if want is not None and field.id != want:
+ raise RuntimeError(
+ "Query auth rules involve column '{}', which the query
projects as a "
+ "different column (field id {} instead of {}); cannot bind
the rules "
+ "by name.".format(field.name, field.id, want))
+ table_field = type_by_name.get(field.name)
+ if table_field is None:
+ continue
+ if field.type != table_field.type:
+ raise RuntimeError(
+ "Query auth rules involve column '{}', which the query
projects with a "
+ "pruned type {} instead of its table type {}; cannot apply
the rules to "
+ "a partial column.".format(field.name, field.type,
table_field.type))
+ # a projected ROW subfield becomes a new top-level column, so the
pruned parent never
+ # shows up by name above
+ complete = {path[0] for path in (nested_name_paths or []) if len(path)
== 1}
+ masking = self.parsed_column_masking()
+ for path in nested_name_paths or []:
+ if len(path) <= 1 or path[0] not in rule_fields:
+ continue
+ if path[0] not in complete:
+ raise RuntimeError(
+ "Query auth rules involve column '{}', which the query
projects only as "
+ "nested field '{}'; cannot apply the rules to a partial
column."
+ .format(path[0], '.'.join(path)))
+ # the whole column is projected too, so the rules still bind to
it; a mask is the
+ # exception, because the nested copy carries the leaf id and stays
raw
+ if path[0] in masking:
+ raise RuntimeError(
+ "Query auth rules mask column '{}', which the query also
projects as "
+ "nested field '{}'; the second copy would be raw."
+ .format(path[0], '.'.join(path)))
+
def get_extra_fields_for_filter(
self,
read_fields: List[DataField],
diff --git a/paimon-python/pypaimon/common/predicate_json_parser.py
b/paimon-python/pypaimon/common/predicate_json_parser.py
index 0f2e5b4e3e..66e79c6881 100644
--- a/paimon-python/pypaimon/common/predicate_json_parser.py
+++ b/paimon-python/pypaimon/common/predicate_json_parser.py
@@ -606,9 +606,8 @@ def _collect_all_field_refs_from_transform(transform: dict,
fields: set = None)
fields.add(transform["fieldRef"]["name"])
else:
for inp in transform.get("inputs", []):
- if isinstance(inp, dict):
- if "name" in inp and "index" in inp:
- fields.add(inp["name"])
- elif "name" in inp:
- _collect_all_field_refs_from_transform(inp, fields)
+ if isinstance(inp, dict) and "name" in inp:
+ # _resolve_transform_input reads every dict input as a column
by name
+ fields.add(inp["name"])
+ _collect_all_field_refs_from_transform(inp, fields)
return fields
diff --git a/paimon-python/pypaimon/read/streaming_table_scan.py
b/paimon-python/pypaimon/read/streaming_table_scan.py
index f34eb48ada..7cd19ee019 100644
--- a/paimon-python/pypaimon/read/streaming_table_scan.py
+++ b/paimon-python/pypaimon/read/streaming_table_scan.py
@@ -36,7 +36,7 @@ from pypaimon.consumer.consumer_manager import ConsumerManager
from pypaimon.manifest.manifest_file_manager import ManifestFileManager
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 wrap_plan_with_auth
from pypaimon.read.scanner.append_table_split_generator import \
AppendTableSplitGenerator
from pypaimon.read.scanner.changelog_follow_up_scanner import \
@@ -275,7 +275,8 @@ class AsyncStreamingTableScan:
return wrap_plan_with_auth(self.__auth_query(), plan)
def __auth_query(self):
- return resolve_auth_result(self._query_auth_fn, self._read_type)
+ from pypaimon.read.table_scan import authorize
+ return authorize(self.table, self._query_auth_fn, self._read_type)
def _start_prefetch(self, snapshot_id: int) -> None:
"""Start prefetching the next scannable snapshot in a background
thread."""
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 61ce079032..c003a828e9 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -184,6 +184,7 @@ class TableRead:
self._parquet_row_group_cache = None
def to_iterator(self, splits: List[Split]) -> Iterator:
+ self._begin_auth_read(splits)
limit = self.limit
def _record_generator():
@@ -793,6 +794,7 @@ class TableRead:
def _arrow_batch_generator(self, splits: List[Split], schema:
pyarrow.Schema,
blob_parallelism: int = 1) ->
Iterator[pyarrow.RecordBatch]:
+ self._begin_auth_read(splits)
chunk_size = 65536
# ``remaining`` tracks how many rows we are still allowed to emit
# across all splits. ``None`` means unlimited.
@@ -963,6 +965,7 @@ class TableRead:
by submission index, so the merged table preserves the order of the
input ``splits`` list.
"""
+ self._begin_auth_read(splits)
remaining_state = _RemainingRows(self.limit)
results: List[Optional[List[pyarrow.RecordBatch]]] = [None] *
len(splits)
workers = min(effective, len(splits))
@@ -1634,6 +1637,15 @@ class TableRead:
"MAP-key projection with query authorization is not supported")
table_fields = self.table.fields
read_fields = self.read_type
+ from pypaimon.read.table_scan import latest_auth_fields,
validate_auth_rules
+ # a split can carry rules without having been planned by a validating
scan
+ snapshot = getattr(self, "_auth_schema", None)
+ if snapshot is None:
+ snapshot = self._auth_schema = latest_auth_fields(self.table)
+ latest_fields = validate_auth_rules(self.table, auth_result, snapshot)
+ # bind to the latest schema: a generated alias may since have become a
real column
+ auth_result.validate_read_type(
+ latest_fields, read_fields, self.nested_name_paths, table_fields)
extra_fields = auth_result.get_extra_fields_for_filter(read_fields,
table_fields)
effective_read_type = read_fields
@@ -1692,6 +1704,21 @@ class TableRead:
return reader
+ def _begin_auth_read(self, splits):
+ """Refresh the schema the rules are validated against, once for this
read.
+
+ Splits are read concurrently and temporal alignment gives each one its
own rule object,
+ so a per-split refresh would ask the catalog once per worker; a
per-reader one would let
+ a later read reuse a schema fetched for an earlier one.
+ """
+ from pypaimon.read.table_scan import latest_auth_fields
+
+ self._auth_schema = None
+ for split in splits:
+ if isinstance(split, QueryAuthSplit) and split.auth_result is not
None:
+ self._auth_schema = latest_auth_fields(self.table)
+ return
+
@staticmethod
def _auth_filter_field_names(auth_result, read_fields) -> set:
filters = getattr(auth_result, "filter", None)
diff --git a/paimon-python/pypaimon/read/table_scan.py
b/paimon-python/pypaimon/read/table_scan.py
index fe1743750c..bcdd203ffc 100755
--- a/paimon-python/pypaimon/read/table_scan.py
+++ b/paimon-python/pypaimon/read/table_scan.py
@@ -367,7 +367,7 @@ class TableScan:
self.file_scanner.skip_level0 = skip_level0
def __auth_query(self):
- return resolve_auth_result(self._query_auth_fn, self._read_type)
+ return authorize(self.table, self._query_auth_fn, self._read_type)
def scan_with_stats(self) -> Tuple[Plan, Optional[ScanStats]]:
"""Run :meth:`plan` while recording manifest / pruning counters.
@@ -630,6 +630,41 @@ class TableScan:
)
+def latest_auth_fields(table):
+ """The table's current columns, re-read so an alter under a long-lived
handle fails closed.
+
+ The catalog answers when it produced the table: a REST data token need not
reach the schema
+ directory, which the scanners avoid for the same reason.
+ """
+ environment = getattr(table, "catalog_environment", None)
+ loader = getattr(environment, "catalog_loader", None)
+ if loader is not None and environment.identifier is not None:
+ return loader.load().get_table(environment.identifier).fields
+ latest = table.schema_manager.latest()
+ return latest.fields if latest is not None else table.fields
+
+
+def validate_auth_rules(table, auth_result, latest_fields=None):
+ """Refuse the read when the rules cannot be bound to the columns they name.
+
+ The reader calls this too, so a split that never met a validating scan is
held to it.
+ """
+ if not auth_result.has_restrictions:
+ return table.fields
+ if latest_fields is None:
+ latest_fields = latest_auth_fields(table)
+ auth_result.validate_against_schema(latest_fields)
+ auth_result.validate_readable_without_rename(latest_fields, table.fields)
+ return latest_fields
+
+
+def authorize(table, query_auth_fn, read_type):
+ auth_result = resolve_auth_result(query_auth_fn, read_type)
+ if auth_result is not None:
+ validate_auth_rules(table, auth_result)
+ return auth_result
+
+
def prune_scanner_by_auth(table, scanner, auth_result):
if not auth_result.filter:
return
diff --git a/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py
b/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py
index 2fbe0fb627..83d47fdc70 100644
--- a/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py
+++ b/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py
@@ -54,8 +54,7 @@ class _BlobCountingFileIO:
return getattr(self._inner, name)
-class _RejectScoreOneAuthResult:
- column_masking = None
+class _RejectScoreOneAuthResult(TableQueryAuthResult):
filter = [json.dumps({
"kind": "LEAF",
"transform": {
@@ -66,6 +65,9 @@ class _RejectScoreOneAuthResult:
"literals": [1],
})]
+ def __init__(self):
+ super().__init__(filter=_RejectScoreOneAuthResult.filter,
column_masking=None)
+
@staticmethod
def get_extra_fields_for_filter(read_fields, table_fields):
return []
@@ -75,12 +77,11 @@ class _RejectScoreOneAuthResult:
return lambda batch: pc.not_equal(batch.column("score"), 1)
-class _PayloadAuthResult:
- column_masking = None
+class _PayloadAuthResult(TableQueryAuthResult):
def __init__(self, expected_payload):
self._expected_payload = expected_payload
- self.filter = [json.dumps({
+ super().__init__(filter=[json.dumps({
"kind": "LEAF",
"transform": {
"name": "FIELD_REF",
@@ -92,7 +93,7 @@ class _PayloadAuthResult:
},
"function": "EQUAL",
"literals": [],
- })]
+ })], column_masking=None)
@staticmethod
def get_extra_fields_for_filter(read_fields, table_fields):
diff --git a/paimon-python/pypaimon/tests/multimodal_temporal_test.py
b/paimon-python/pypaimon/tests/multimodal_temporal_test.py
index 6ffbc80368..c34acf0197 100644
--- a/paimon-python/pypaimon/tests/multimodal_temporal_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py
@@ -1970,7 +1970,7 @@ class MultimodalTemporalTest(unittest.TestCase):
for select in selected
))
- def test_alignment_does_not_mask_internal_key_dependencies(self):
+ def test_alignment_rejects_a_mask_reading_a_masked_dependency(self):
anchors = self._table("dependency_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
@@ -2014,19 +2014,14 @@ class MultimodalTemporalTest(unittest.TestCase):
anchors.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
- result = pmm.join_asof(
- anchors.scan(), source.scan().select("value"),
- on="event_time", by="episode_id",
- direction="nearest", tolerance=0,
- ).to_arrow()
-
- self.assertEqual(7, result["value"][0].as_py())
- self.assertEqual(1, result["episode_id"][0].as_py())
- self.assertEqual("1", result["allowed_episode_id"][0].as_py())
- self.assertEqual(
- pa.string(), result.schema.field("allowed_episode_id").type)
+ with self.assertRaisesRegex(ValueError, "which is masked too"):
+ pmm.join_asof(
+ anchors.scan(), source.scan().select("value"),
+ on="event_time", by="episode_id",
+ direction="nearest", tolerance=0,
+ ).to_arrow()
- def test_alignment_matches_masking_reader_rule_semantics(self):
+ def test_alignment_rejects_cross_column_and_json_null_masks(self):
anchors = self._table("mask_semantics_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
@@ -2064,21 +2059,21 @@ class MultimodalTemporalTest(unittest.TestCase):
source.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth[0])
- row = pmm.join_asof(
- anchors.scan(), source.scan().select(["first", "second"]),
- on="event_time", by="episode_id",
- direction="nearest", tolerance=0,
- ).to_list()[0]
- self.assertEqual(("b", "a"), (row["first"], row["second"]))
+ with self.assertRaisesRegex(ValueError, "which is masked too"):
+ pmm.join_asof(
+ anchors.scan(), source.scan().select(["first", "second"]),
+ on="event_time", by="episode_id",
+ direction="nearest", tolerance=0,
+ ).to_list()
auth[0] = TableQueryAuthResult(
filter=None, column_masking={"first": "null"})
- row = pmm.join_asof(
- anchors.scan(), source.scan().select("first"),
- on="event_time", by="episode_id",
- direction="nearest", tolerance=0,
- ).to_list()[0]
- self.assertEqual("a", row["first"])
+ with self.assertRaisesRegex(ValueError, "JSON null"):
+ pmm.join_asof(
+ anchors.scan(), source.scan().select("first"),
+ on="event_time", by="episode_id",
+ direction="nearest", tolerance=0,
+ ).to_list()
def test_alignment_rejects_incremental_scans(self):
anchors = self.conn.create_table(
@@ -2282,7 +2277,7 @@ class MultimodalTemporalTest(unittest.TestCase):
source.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
- with self.assertRaisesRegex(ValueError, "nested projection"):
+ with self.assertRaisesRegex(ValueError, "does not exist in table
schema"):
pmm.join_asof(
anchors.scan(), source.scan().select("payload.value"),
on="event_time", by="episode_id",
diff --git a/paimon-python/pypaimon/tests/query_auth_validation_test.py
b/paimon-python/pypaimon/tests/query_auth_validation_test.py
new file mode 100644
index 0000000000..d95affffab
--- /dev/null
+++ b/paimon-python/pypaimon/tests/query_auth_validation_test.py
@@ -0,0 +1,957 @@
+################################################################################
+# 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.
+################################################################################
+
+import json
+import time
+import unittest
+
+from types import SimpleNamespace
+
+import pyarrow as pa
+
+from pypaimon.common.options import CoreOptions, Options
+from pypaimon.read.query_auth_split import QueryAuthSplit
+from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
+from pypaimon.read.table_read import TableRead
+from pypaimon.schema.data_types import PyarrowFieldParser
+from pypaimon.catalog.table_query_auth import TableQueryAuthResult
+from pypaimon.read.table_scan import TableScan, validate_auth_rules
+from pypaimon.read.streaming_table_scan import AsyncStreamingTableScan
+from pypaimon.read.plan import Plan
+from pypaimon.schema.data_types import AtomicType, DataField, RowType
+from pypaimon.table.special_fields import SpecialFields
+
+
+def _string_field(field_id, name):
+ return DataField(field_id, name, AtomicType("STRING"))
+
+
+def _struct_field(field_id, name, children):
+ return DataField(field_id, name, RowType(True, children))
+
+
+def _null_mask():
+ return json.dumps({"name": "NULL"})
+
+
+def _field_ref_mask(source, index=0):
+ return json.dumps({
+ "name": "FIELD_REF",
+ "fieldRef": {"index": index, "name": source, "type": "STRING"},
+ })
+
+
+def _filter_json(field_name="dept", value="eng"):
+ return json.dumps({
+ "kind": "LEAF",
+ "transform": {
+ "name": "FIELD_REF",
+ "fieldRef": {"index": 0, "name": field_name, "type": "STRING"},
+ },
+ "function": "EQUAL",
+ "literals": [value],
+ })
+
+
+class _FakeSchema:
+ def __init__(self, fields):
+ self.fields = fields
+
+
+class _FakeSchemaManager:
+ def __init__(self, fields, error=None):
+ self._fields = fields
+ self._error = error
+
+ def latest(self):
+ if self._error is not None:
+ raise self._error
+ return _FakeSchema(self._fields)
+
+
+class _FakeOptions:
+ def blob_descriptor_fields(self):
+ return set()
+
+ def blob_view_fields(self):
+ return set()
+
+ def data_evolution_enabled(self):
+ return False
+
+ def row_tracking_enabled(self):
+ return False
+
+
+class _FakeTable:
+ def __init__(self, fields, latest_fields=None, latest_error=None):
+ self.fields = fields
+ self.options = _FakeOptions()
+ self.schema_manager = _FakeSchemaManager(
+ fields if latest_fields is None else latest_fields, latest_error)
+
+
+def _make_scan(auth_result, read_type, table):
+ scan = TableScan.__new__(TableScan)
+ scan._query_auth_fn = lambda select: auth_result
+ scan._read_type = read_type
+ scan.table = table
+ return scan
+
+
+def _make_read(table, read_type, nested_name_paths=None):
+ read = TableRead.__new__(TableRead)
+ read.table = table
+ read.read_type = read_type
+ read.nested_name_paths = nested_name_paths
+ return read
+
+
+class TestValidateAgainstSchema(unittest.TestCase):
+
+ def test_stale_mask_after_rename_is_rejected(self):
+ result = TableQueryAuthResult(None, {"email": _null_mask()})
+ fields = [_string_field(0, "id"), _string_field(1, "mail")]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("email", str(ctx.exception))
+ self.assertIn("rename or drop", str(ctx.exception))
+
+ def test_stale_filter_after_rename_is_rejected(self):
+ result = TableQueryAuthResult([_filter_json("dept")], None)
+ fields = [_string_field(0, "id"), _string_field(1, "team")]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("Row filter", str(ctx.exception))
+ self.assertIn("dept", str(ctx.exception))
+
+ def test_missing_mask_column_names_the_column(self):
+ result = TableQueryAuthResult(None, {"gone": _null_mask()})
+ fields = [_string_field(0, "id")]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("gone", str(ctx.exception))
+ self.assertIn("does not exist in table schema", str(ctx.exception))
+
+ def test_missing_filter_column_names_the_column(self):
+ result = TableQueryAuthResult([_filter_json("gone")], None)
+ fields = [_string_field(0, "id")]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("gone", str(ctx.exception))
+ self.assertIn("does not exist in table schema", str(ctx.exception))
+
+ def test_mask_reading_a_masked_column_is_rejected(self):
+ result = TableQueryAuthResult(None, {
+ "email": _null_mask(),
+ "alias": _field_ref_mask("email", 1),
+ })
+ fields = [
+ _string_field(0, "id"),
+ _string_field(1, "email"),
+ _string_field(2, "alias"),
+ ]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ message = str(ctx.exception)
+ self.assertIn("alias", message)
+ self.assertIn("email", message)
+ self.assertIn("raw value", message)
+
+ def test_mask_reading_itself_is_allowed(self):
+ result = TableQueryAuthResult(None, {"email": _field_ref_mask("email",
1)})
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result.validate_against_schema(fields)
+
+ def test_mask_reading_an_unmasked_column_is_allowed(self):
+ result = TableQueryAuthResult(None, {"alias": _field_ref_mask("email",
1)})
+ fields = [
+ _string_field(0, "id"),
+ _string_field(1, "email"),
+ _string_field(2, "alias"),
+ ]
+ result.validate_against_schema(fields)
+
+ def test_a_system_column_is_never_missing_from_the_schema(self):
+ fields = [_string_field(0, "id")]
+ for rules in (
+ (None, {SpecialFields.ROW_ID.name: _null_mask()}),
+ (None, {"id": _field_ref_mask("_SEQUENCE_NUMBER", 1)}),
+ ([_filter_json("_VALUE_KIND", "x")], None)):
+ TableQueryAuthResult(*rules).validate_against_schema(fields)
+
+ def test_valid_rules_pass(self):
+ result = TableQueryAuthResult([_filter_json("dept")], {"email":
_null_mask()})
+ fields = [
+ _string_field(0, "id"),
+ _string_field(1, "email"),
+ _string_field(2, "dept"),
+ ]
+ result.validate_against_schema(fields)
+
+ def test_no_rules_pass(self):
+ result = TableQueryAuthResult(None, None)
+ result.validate_against_schema([_string_field(0, "id")])
+
+
+class TestValidateReadableWithoutRename(unittest.TestCase):
+
+ def test_renamed_column_is_rejected(self):
+ result = TableQueryAuthResult(None, {"mail": _null_mask()})
+ latest = [_string_field(0, "id"), _string_field(1, "mail")]
+ read = [_string_field(0, "id"), _string_field(1, "email")]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_readable_without_rename(latest, read)
+ message = str(ctx.exception)
+ self.assertIn("mail", message)
+ self.assertIn("email", message)
+ self.assertIn("renamed since", message)
+
+ def test_renamed_filter_operand_is_rejected(self):
+ result = TableQueryAuthResult([_filter_json("team")], None)
+ latest = [_string_field(0, "id"), _string_field(1, "team")]
+ read = [_string_field(0, "id"), _string_field(1, "dept")]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_readable_without_rename(latest, read)
+ self.assertIn("Row filter", str(ctx.exception))
+ self.assertIn("renamed since", str(ctx.exception))
+
+ def test_dropped_and_readded_column_is_rejected(self):
+ result = TableQueryAuthResult(None, {"email": _null_mask()})
+ latest = [_string_field(0, "id"), _string_field(7, "email")]
+ read = [_string_field(0, "id"), _string_field(1, "email")]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_readable_without_rename(latest, read)
+ self.assertIn("dropped and re-added", str(ctx.exception))
+
+ def test_matching_ids_pass(self):
+ result = TableQueryAuthResult(None, {"email": _null_mask()})
+ latest = [_string_field(0, "id"), _string_field(1, "email")]
+ result.validate_readable_without_rename(latest, list(latest))
+
+ def test_column_absent_from_read_schema_is_inert(self):
+ result = TableQueryAuthResult(None, {"email": _null_mask()})
+ latest = [_string_field(0, "id"), _string_field(1, "email")]
+ read = [_string_field(0, "id")]
+ result.validate_readable_without_rename(latest, read)
+
+ def test_mask_input_renamed_is_rejected_when_target_readable(self):
+ result = TableQueryAuthResult(None, {"alias": _field_ref_mask("email",
1)})
+ latest = [
+ _string_field(0, "id"),
+ _string_field(1, "email"),
+ _string_field(2, "alias"),
+ ]
+ read = [
+ _string_field(0, "id"),
+ _string_field(1, "mail"),
+ _string_field(2, "alias"),
+ ]
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_readable_without_rename(latest, read)
+ self.assertIn("renamed since", str(ctx.exception))
+
+ def test_column_absent_from_latest_schema_is_skipped(self):
+ result = TableQueryAuthResult(None, {"gone": _null_mask()})
+ latest = [_string_field(0, "id")]
+ result.validate_readable_without_rename(latest, list(latest))
+
+
+class TestValidateReadType(unittest.TestCase):
+
+ def _table_fields(self):
+ return [
+ _string_field(0, "id"),
+ _struct_field(1, "s", [_string_field(2, "a"), _string_field(3,
"b")]),
+ _string_field(4, "dept"),
+ ]
+
+ def test_nested_subfield_projection_of_mask_target_is_rejected(self):
+ result = TableQueryAuthResult(None, {"s": _null_mask()})
+ read = [_string_field(0, "id"), _string_field(2, "s_a")]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(
+ self._table_fields(), read, [["id"], ["s", "a"]])
+ message = str(ctx.exception)
+ self.assertIn("'s'", message)
+ self.assertIn("s.a", message)
+ self.assertIn("partial column", message)
+
+ def
test_nested_subfield_projection_of_inactive_mask_input_is_allowed(self):
+ result = TableQueryAuthResult(None, {"dept": _field_ref_mask("s", 1)})
+ read = [_string_field(0, "id"), _string_field(2, "s_a")]
+ result.validate_read_type(self._table_fields(), read, [["id"], ["s",
"a"]])
+
+ def test_nested_subfield_projection_of_filter_operand_is_rejected(self):
+ result = TableQueryAuthResult([_filter_json("s")], None)
+ read = [_string_field(0, "id"), _string_field(2, "s_a")]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(
+ self._table_fields(), read, [["id"], ["s", "a"]])
+ self.assertIn("partial column", str(ctx.exception))
+
+ def test_a_column_flattened_onto_a_system_name_is_rejected(self):
+ table_fields = [
+ _string_field(0, "id"),
+ _struct_field(1, "_ROW", [_string_field(2, "ID")]),
+ ]
+ result = TableQueryAuthResult([_filter_json("_ROW_ID", "999")], None)
+ read = [_string_field(0, "id"), _string_field(2, "_ROW_ID")]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(table_fields, read, [["id"], ["_ROW",
"ID"]])
+ self.assertIn("field id", str(ctx.exception))
+
+ def
test_a_duplicate_projection_of_a_masked_system_column_is_rejected(self):
+ row_id = SpecialFields.ROW_ID
+ result = TableQueryAuthResult(None, {row_id.name: _null_mask()})
+ read = [
+ DataField(row_id.id, row_id.name, AtomicType("BIGINT")),
+ DataField(row_id.id, row_id.name + "__0", AtomicType("BIGINT")),
+ ]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(self._table_fields(), read)
+ self.assertIn("would be raw", str(ctx.exception))
+
+ def test_a_single_projection_of_a_masked_system_column_is_allowed(self):
+ row_id = SpecialFields.ROW_ID
+ result = TableQueryAuthResult(None, {row_id.name: _null_mask()})
+ read = [_string_field(0, "id"),
+ DataField(row_id.id, row_id.name, AtomicType("BIGINT"))]
+ result.validate_read_type(self._table_fields(), read)
+
+ def
test_nested_subfield_beside_the_complete_filter_operand_is_allowed(self):
+ result = TableQueryAuthResult([_filter_json("s")], None)
+ read = [
+ _string_field(0, "id"),
+ _struct_field(1, "s", [_string_field(2, "a"), _string_field(3,
"b")]),
+ _string_field(2, "s_a"),
+ ]
+ result.validate_read_type(
+ self._table_fields(), read, [["id"], ["s"], ["s", "a"]])
+
+ def test_nested_subfield_beside_the_complete_mask_target_is_rejected(self):
+ result = TableQueryAuthResult(None, {"s": _null_mask()})
+ read = [
+ _string_field(0, "id"),
+ _struct_field(1, "s", [_string_field(2, "a"), _string_field(3,
"b")]),
+ _string_field(2, "s_a"),
+ ]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(
+ self._table_fields(), read, [["id"], ["s"], ["s", "a"]])
+ self.assertIn("would be raw", str(ctx.exception))
+
+ def test_nested_subfield_projection_of_unrelated_column_is_allowed(self):
+ result = TableQueryAuthResult([_filter_json("dept")], {"id":
_null_mask()})
+ read = [
+ _string_field(0, "id"),
+ _string_field(4, "dept"),
+ _string_field(2, "s_a"),
+ ]
+ result.validate_read_type(
+ self._table_fields(), read, [["id"], ["dept"], ["s", "a"]])
+
+ def test_top_level_projection_of_mask_target_is_allowed(self):
+ result = TableQueryAuthResult(None, {"s": _null_mask()})
+ table_fields = self._table_fields()
+ read = [table_fields[0], table_fields[1]]
+ result.validate_read_type(table_fields, read, None)
+
+ def test_pruned_type_in_read_schema_is_rejected(self):
+ result = TableQueryAuthResult(None, {"s": _null_mask()})
+ table_fields = self._table_fields()
+ read = [
+ _string_field(0, "id"),
+ _struct_field(1, "s", [_string_field(2, "a")]),
+ ]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(table_fields, read, None)
+ self.assertIn("pruned type", str(ctx.exception))
+
+ def test_no_rules_pass(self):
+ result = TableQueryAuthResult(None, None)
+ result.validate_read_type(
+ self._table_fields(), [_string_field(2, "s_a")], [["s", "a"]])
+
+
+class TestScanValidationWiring(unittest.TestCase):
+
+ def test_scan_rejects_stale_mask_after_rename(self):
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ fields = [_string_field(0, "id"), _string_field(1, "mail")]
+ scan = _make_scan(auth, fields, _FakeTable(fields))
+ with self.assertRaises(ValueError) as ctx:
+ scan._TableScan__auth_query()
+ self.assertIn("email", str(ctx.exception))
+
+ def test_scan_rejects_mask_reading_a_masked_column(self):
+ auth = TableQueryAuthResult(None, {
+ "email": _null_mask(),
+ "alias": _field_ref_mask("email", 1),
+ })
+ fields = [
+ _string_field(0, "id"),
+ _string_field(1, "email"),
+ _string_field(2, "alias"),
+ ]
+ scan = _make_scan(auth, fields, _FakeTable(fields))
+ with self.assertRaises(ValueError) as ctx:
+ scan._TableScan__auth_query()
+ self.assertIn("raw value", str(ctx.exception))
+
+ def test_scan_rejects_mask_stale_against_latest_schema_only(self):
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ read_fields = [_string_field(0, "id"), _string_field(1, "email")]
+ latest_fields = [_string_field(0, "id"), _string_field(1, "mail")]
+ scan = _make_scan(
+ auth, read_fields, _FakeTable(read_fields,
latest_fields=latest_fields))
+ with self.assertRaises(ValueError) as ctx:
+ scan._TableScan__auth_query()
+ self.assertIn("email", str(ctx.exception))
+
+ def test_scan_rejects_rule_column_renamed_in_read_schema(self):
+ auth = TableQueryAuthResult(None, {"mail": _null_mask()})
+ read_fields = [_string_field(0, "id"), _string_field(1, "email")]
+ latest_fields = [_string_field(0, "id"), _string_field(1, "mail")]
+ scan = _make_scan(
+ auth, read_fields, _FakeTable(read_fields,
latest_fields=latest_fields))
+ with self.assertRaises(ValueError) as ctx:
+ scan._TableScan__auth_query()
+ self.assertIn("renamed since", str(ctx.exception))
+
+ def test_scan_allows_valid_rules(self):
+ auth = TableQueryAuthResult([_filter_json("dept")], {"email":
_null_mask()})
+ fields = [
+ _string_field(0, "id"),
+ _string_field(1, "email"),
+ _string_field(2, "dept"),
+ ]
+ scan = _make_scan(auth, fields, _FakeTable(fields))
+ self.assertIs(scan._TableScan__auth_query(), auth)
+
+ def test_scan_without_rules_skips_validation(self):
+ fields = [_string_field(0, "id")]
+ scan = _make_scan(TableQueryAuthResult(None, None), fields,
_FakeTable(fields))
+ self.assertIsNone(scan._TableScan__auth_query())
+
+ def test_scan_propagates_a_latest_schema_read_failure(self):
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ table = _FakeTable(fields, latest_error=RuntimeError("no schema"))
+ scan = _make_scan(auth, fields, table)
+ with self.assertRaises(RuntimeError):
+ scan._TableScan__auth_query()
+
+
+class TestReadValidationWiring(unittest.TestCase):
+
+ def _table_fields(self):
+ return [
+ _string_field(0, "id"),
+ _struct_field(1, "s", [_string_field(2, "a"), _string_field(3,
"b")]),
+ ]
+
+ def test_read_rejects_nested_projection_of_mask_target(self):
+ auth = TableQueryAuthResult(None, {"s": _null_mask()})
+ table_fields = self._table_fields()
+ read = _make_read(
+ _FakeTable(table_fields),
+ [_string_field(0, "id"), _string_field(2, "s_a")],
+ [["id"], ["s", "a"]])
+ with self.assertRaises(RuntimeError) as ctx:
+ read._TableRead__authed_reader(None, auth)
+ self.assertIn("partial column", str(ctx.exception))
+
+ def test_read_rejects_nested_projection_of_filter_operand(self):
+ auth = TableQueryAuthResult([_filter_json("s")], None)
+ table_fields = self._table_fields()
+ read = _make_read(
+ _FakeTable(table_fields),
+ [_string_field(0, "id"), _string_field(2, "s_a")],
+ [["id"], ["s", "a"]])
+ with self.assertRaises(RuntimeError) as ctx:
+ read._TableRead__authed_reader(None, auth)
+ self.assertIn("partial column", str(ctx.exception))
+
+
+if __name__ == '__main__':
+ unittest.main()
+
+
+class TestCodexReviewFindings(unittest.TestCase):
+
+ def
test_projection_binding_a_rule_name_to_another_column_is_rejected(self):
+ table_fields = [_string_field(0, "id"), _string_field(1, "s_a"),
+ _struct_field(2, "s", [_string_field(3, "a")])]
+ result = TableQueryAuthResult(None, {"s_a": _null_mask()})
+ read = [_string_field(0, "id"), _string_field(3, "s_a")]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(table_fields, read, [["id"], ["s", "a"]])
+ self.assertIn("field id", str(ctx.exception))
+
+ def test_json_null_mask_is_rejected(self):
+ result = TableQueryAuthResult(None, {"email": json.dumps(None)})
+ with self.assertRaises(ValueError) as ctx:
+ result.parsed_column_masking()
+ self.assertIn("JSON null", str(ctx.exception))
+
+ def test_inactive_mask_inputs_are_not_rule_fields(self):
+ result = TableQueryAuthResult(None, {"dept": _field_ref_mask("s", 1)})
+ self.assertNotIn("s", result.rule_field_names({"id"}))
+ self.assertIn("s", result.rule_field_names({"dept"}))
+
+ def
test_a_mask_on_a_renamed_physical_system_looking_column_is_rejected(self):
+ fields = [_string_field(0, "id"), _string_field(1, "renamed")]
+ result = TableQueryAuthResult(None, {"rowkind": _null_mask()})
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("rowkind", str(ctx.exception))
+
+ def
test_a_cross_mask_on_a_physical_system_looking_column_is_rejected(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email"),
+ _string_field(2, "_LEVEL")]
+ result = TableQueryAuthResult(None, {
+ "email": _null_mask(), "_LEVEL": _field_ref_mask("email", 1)})
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("masked too", str(ctx.exception))
+
+ def test_streaming_scan_validates_like_the_batch_scan(self):
+ from pypaimon.read.streaming_table_scan import AsyncStreamingTableScan
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ fields = [_string_field(0, "id"), _string_field(1, "mail")]
+ scan = AsyncStreamingTableScan.__new__(AsyncStreamingTableScan)
+ scan._query_auth_fn = lambda select: auth
+ scan._read_type = fields
+ scan.table = _FakeTable(fields)
+ with self.assertRaises(ValueError) as ctx:
+ scan._AsyncStreamingTableScan__auth_query()
+ self.assertIn("does not exist", str(ctx.exception))
+
+ def test_duplicate_projection_alias_of_a_masked_column_is_rejected(self):
+ table_fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result = TableQueryAuthResult(None, {"email": _null_mask()})
+ read = [_string_field(1, "email"), _string_field(1, "email__0")]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(table_fields, read, None)
+ self.assertIn("email__0", str(ctx.exception))
+
+ def test_ordinary_mask_input_must_exist(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result = TableQueryAuthResult(None, {"email": _field_ref_mask("gone",
2)})
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("'gone'", str(ctx.exception))
+
+ def test_inputs_of_a_mask_absent_from_the_read_schema_are_inert(self):
+ latest = [_string_field(0, "id"), _string_field(1, "email")]
+ read = [_string_field(0, "id"), _string_field(7, "email")]
+ result = TableQueryAuthResult(None, {"gone": _field_ref_mask("email",
1)})
+ result.validate_readable_without_rename(latest, read)
+
+ def test_pruned_type_is_checked_for_filter_operands_too(self):
+ table_fields = [_string_field(0, "id"),
+ _struct_field(1, "s", [_string_field(2, "a"),
_string_field(3, "b")])]
+ result = TableQueryAuthResult([_filter_json("s", "x")], None)
+ read = [_string_field(0, "id"), _struct_field(1, "s",
[_string_field(2, "a")])]
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(table_fields, read, None)
+ self.assertIn("pruned type", str(ctx.exception))
+
+ def test_duplicate_alias_of_an_unmasked_filter_operand_is_allowed(self):
+ table_fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result = TableQueryAuthResult([_filter_json("email", "x")], None)
+ read = [_string_field(1, "email"), _string_field(1, "email__0")]
+ result.validate_read_type(table_fields, read, None)
+
+ def test_duplicate_alias_of_an_unmasked_mask_input_is_allowed(self):
+ table_fields = [_string_field(0, "id"), _string_field(1, "email"),
+ _string_field(2, "alias")]
+ result = TableQueryAuthResult(None, {"alias": _field_ref_mask("email",
1)})
+ read = [_string_field(2, "alias"), _string_field(1, "email"),
+ _string_field(1, "email__0")]
+ result.validate_read_type(table_fields, read, None)
+
+
+class TestStreamingPlanAuthorization(unittest.TestCase):
+
+ def _scan(self, stale):
+ fields = [_string_field(0, "email")]
+ latest = [_string_field(0, "mail")] if stale else fields
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ scan = AsyncStreamingTableScan.__new__(AsyncStreamingTableScan)
+ scan.table = _FakeTable(fields, latest_fields=latest)
+ scan._read_type = fields
+ scan._query_auth_fn = lambda select: auth
+ raw = Plan([object()], snapshot_id=7)
+ scan._AsyncStreamingTableScan__create_initial_plan_raw = lambda *a: raw
+ scan._create_delta_plan = lambda *a: raw
+ scan._create_changelog_plan = lambda *a: raw
+ scan.follow_up_scanner = object()
+ return scan, auth
+
+ def _entry(self, scan, path):
+ if path == "initial":
+ return lambda: scan._create_initial_plan(None)
+ if path == "catchup":
+ return lambda: scan._create_catch_up_plan(1, None)
+ return lambda: scan._create_follow_up_plan(None)
+
+ def test_every_plan_path_wraps_with_the_authorization(self):
+ for path in ("initial", "delta", "catchup"):
+ scan, auth = self._scan(stale=False)
+ plan = self._entry(scan, path)()
+ self.assertIsInstance(plan.splits()[0], QueryAuthSplit)
+ self.assertIs(plan.splits()[0].auth_result, auth)
+
+ def test_every_plan_path_refuses_a_stale_rule(self):
+ for path in ("initial", "delta", "catchup"):
+ scan, _ = self._scan(stale=True)
+ with self.assertRaises(ValueError):
+ self._entry(scan, path)()
+
+ def
test_a_system_column_mask_survives_a_rename_of_every_other_column(self):
+ result = TableQueryAuthResult(None, {"_VALUE_KIND": _null_mask()})
+ fields = [_string_field(0, "id"), _string_field(1, "renamed")]
+ result.validate_against_schema(fields)
+
+ def test_duplicate_alias_of_a_column_outside_the_rules_is_allowed(self):
+ table_fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result = TableQueryAuthResult(None, {"email": _null_mask()})
+ read = [_string_field(1, "email"), _string_field(0, "id"),
+ _string_field(0, "id__0")]
+ result.validate_read_type(table_fields, read, None)
+
+ def test_an_indexless_transform_input_is_discovered(self):
+ mask = json.dumps({"name": "CONCAT",
+ "inputs": [{"name": "gone", "type": "STRING"}]})
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result = TableQueryAuthResult(None, {"email": mask})
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("'gone'", str(ctx.exception))
+
+ def test_an_indexless_input_reading_a_masked_column_is_rejected(self):
+ mask = json.dumps({"name": "CONCAT",
+ "inputs": [{"name": "email", "type": "STRING"}]})
+ fields = [_string_field(0, "email"), _string_field(1, "alias")]
+ result = TableQueryAuthResult(None, {"email": _null_mask(), "alias":
mask})
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("masked too", str(ctx.exception))
+
+ def test_an_indexless_system_input_is_refused_when_unprojected(self):
+ for name in ("_ROW_ID", "_SEQUENCE_NUMBER", "_VALUE_KIND"):
+ mask = json.dumps({"name": "CONCAT",
+ "inputs": [{"name": name, "type": "STRING"}]})
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result = TableQueryAuthResult(None, {"email": mask})
+ result.validate_against_schema(fields)
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(fields, fields)
+ self.assertIn(name, str(ctx.exception))
+
+ def
test_a_structural_input_carrying_both_index_and_inputs_is_discovered(self):
+ mask = json.dumps({"name": "CONCAT", "inputs": [
+ {"name": "email", "index": 0, "type": "STRING", "inputs": []}]})
+ fields = [_string_field(0, "email"), _string_field(1, "alias")]
+ result = TableQueryAuthResult(None, {"email": _null_mask(), "alias":
mask})
+ with self.assertRaises(ValueError) as ctx:
+ result.validate_against_schema(fields)
+ self.assertIn("masked too", str(ctx.exception))
+
+ def test_a_structural_system_input_is_refused_when_unprojected(self):
+ mask = json.dumps({"name": "CONCAT", "inputs": [
+ {"name": "_ROW_ID", "index": 0, "type": "STRING", "inputs": []}]})
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ result = TableQueryAuthResult(None, {"email": mask})
+ result.validate_against_schema(fields)
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(fields, fields)
+ self.assertIn("_ROW_ID", str(ctx.exception))
+
+ def test_a_system_input_is_allowed_once_the_query_projects_it(self):
+ mask = _field_ref_mask("_ROW_ID", 1)
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ read = fields + [_string_field(2147483642, "_ROW_ID")]
+ result = TableQueryAuthResult(None, {"email": mask})
+ result.validate_read_type(fields, read)
+
+ def test_a_mask_on_an_unprojected_system_column_stays_inert(self):
+ fields = [_string_field(0, "id")]
+ result = TableQueryAuthResult(None, {"_ROW_ID": _null_mask()})
+ result.validate_read_type(fields, fields)
+
+ def
test_a_rule_naming_a_column_added_after_a_generated_alias_is_rejected(self):
+ latest = [_string_field(0, "email"), _string_field(3, "email__0")]
+ read = [_string_field(0, "email"), _string_field(0, "email__0")]
+ snapshot = [_string_field(0, "email")]
+ result = TableQueryAuthResult(None, {"email__0": _null_mask()})
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(latest, read, None, snapshot)
+ self.assertIn("field id", str(ctx.exception))
+
+ def
test_a_filter_naming_a_column_added_after_a_generated_alias_is_rejected(self):
+ latest = [_string_field(0, "email"), _string_field(3, "s_a")]
+ read = [_string_field(0, "email"), _string_field(2, "s_a")]
+ snapshot = [_string_field(0, "email"), _string_field(2, "s_a")]
+ result = TableQueryAuthResult([_filter_json("s_a", "eng")], None)
+ with self.assertRaises(RuntimeError):
+ result.validate_read_type(latest, read, None, snapshot)
+
+ def test_a_widened_column_type_is_still_readable(self):
+ latest = [DataField(0, "n", AtomicType("BIGINT"))]
+ snapshot = [DataField(0, "n", AtomicType("INT"))]
+ result = TableQueryAuthResult(None, {"n": _null_mask()})
+ result.validate_read_type(latest, snapshot, None, snapshot)
+
+ def test_a_pruned_type_is_still_rejected_against_the_snapshot(self):
+ latest = [_string_field(0, "id"),
+ _struct_field(1, "s", [_string_field(2, "a"),
_string_field(3, "b")])]
+ snapshot = latest
+ read = [_string_field(0, "id"), _struct_field(1, "s",
[_string_field(2, "a")])]
+ result = TableQueryAuthResult(None, {"s": _null_mask()})
+ with self.assertRaises(RuntimeError) as ctx:
+ result.validate_read_type(latest, read, None, snapshot)
+ self.assertIn("pruned type", str(ctx.exception))
+
+ def test_an_unrestricted_result_needs_no_latest_schema(self):
+ from pypaimon.read.table_scan import validate_auth_rules
+ table = _FakeTable([_string_field(0, "id")],
+ latest_error=RuntimeError("must not be read"))
+ validate_auth_rules(table, TableQueryAuthResult(None, None))
+
+ def test_a_filter_only_result_is_still_validated(self):
+ from pypaimon.read.table_scan import validate_auth_rules
+ table = _FakeTable([_string_field(0, "id")])
+ auth = TableQueryAuthResult([_filter_json("gone", "x")], None)
+ with self.assertRaises(ValueError):
+ validate_auth_rules(table, auth)
+
+
+class _OneBatch(RecordBatchReader):
+ def __init__(self, batch):
+ self._batch = batch
+
+ def read_arrow_batch(self):
+ batch, self._batch = self._batch, None
+ return batch
+
+ def close(self):
+ pass
+
+
+class _CountingCatalog:
+ def __init__(self, fields):
+ self.identifier = "db.t"
+ self.catalog_loader = self
+ self._fields = fields
+ self.get_table_calls = 0
+
+ def load(self):
+ return self
+
+ def get_table(self, identifier):
+ self.get_table_calls += 1
+ time.sleep(0.05)
+ return SimpleNamespace(fields=self._fields)
+
+
+def _exploding_schema_manager():
+ def boom():
+ raise AssertionError("the schema directory must not be listed")
+ return SimpleNamespace(latest=boom)
+
+
+class TestSchemaRefreshAsksTheCatalog(unittest.TestCase):
+
+ def _table(self, cached, latest):
+ return SimpleNamespace(
+ fields=cached, options=CoreOptions(Options({})),
+ is_primary_key_table=False,
+ schema_manager=_exploding_schema_manager(),
+ catalog_environment=_CountingCatalog(latest))
+
+ def test_a_catalog_backed_table_never_lists_the_schema_directory(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ table = self._table(fields, fields)
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ self.assertEqual(
+ [f.name for f in validate_auth_rules(table, auth)], ["id",
"email"])
+ self.assertEqual(table.catalog_environment.get_table_calls, 1)
+
+ def test_a_rename_seen_only_by_the_catalog_still_fails_closed(self):
+ cached = [_string_field(0, "id"), _string_field(1, "email")]
+ renamed = [_string_field(0, "id"), _string_field(1, "mail")]
+ table = self._table(cached, renamed)
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ with self.assertRaises(ValueError):
+ validate_auth_rules(table, auth)
+
+ def _reader(self, table, fields):
+ read = TableRead(table, None, fields)
+ batch = pa.RecordBatch.from_pydict(
+ {"id": ["1"], "email": ["SECRET"]},
+ schema=PyarrowFieldParser.from_paimon_schema(fields))
+ read._create_split_read = lambda split, **kw: SimpleNamespace(
+ create_reader=lambda: _OneBatch(batch))
+ return read
+
+ def
test_a_read_asks_the_catalog_once_across_workers_and_rule_objects(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ table = self._table(fields, fields)
+ read = self._reader(table, fields)
+ splits = [
+ QueryAuthSplit(object(), TableQueryAuthResult(None, {"email":
_null_mask()}))
+ for _ in range(8)]
+ self.assertEqual(
+ read.to_arrow(splits, parallelism=4).to_pydict(),
+ {"id": ["1"] * 8, "email": [None] * 8})
+ self.assertEqual(table.catalog_environment.get_table_calls, 1)
+
+ def test_a_reused_reader_refreshes_between_reads(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ table = self._table(fields, fields)
+ read = self._reader(table, fields)
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ for _ in range(3):
+ read.to_arrow([QueryAuthSplit(object(), auth)], parallelism=1)
+ self.assertEqual(table.catalog_environment.get_table_calls, 3)
+
+ def test_a_read_without_rules_never_asks_the_catalog(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ table = self._table(fields, fields)
+ read = self._reader(table, fields)
+ read.to_arrow([object()], parallelism=1)
+ self.assertEqual(table.catalog_environment.get_table_calls, 0)
+
+ def test_the_reader_stays_picklable_after_a_read(self):
+ import pickle
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ table = self._table(fields, fields)
+ read = self._reader(table, fields)
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ read.to_arrow([QueryAuthSplit(object(), auth)], parallelism=1)
+ pickle.dumps(read._auth_schema)
+
+
+class TestReaderAppliesTheRules(unittest.TestCase):
+
+ def _read(self, fields, data, latest=None, projected=None):
+ latest_schema = SimpleNamespace(fields=fields if latest is None else
latest)
+ table = SimpleNamespace(
+ fields=fields, options=CoreOptions(Options({})),
is_primary_key_table=False,
+ schema_manager=SimpleNamespace(latest=lambda: latest_schema))
+ read = TableRead(table, None, fields if projected is None else
projected)
+
+ def storage(split, **kwargs):
+ requested = kwargs.get("read_type", read.read_type)
+ schema = PyarrowFieldParser.from_paimon_schema(requested)
+ batch = pa.RecordBatch.from_pydict(
+ {f.name: data[f.name] for f in requested}, schema=schema)
+ return SimpleNamespace(create_reader=lambda: _OneBatch(batch))
+
+ read._create_split_read = storage
+ return read
+
+ def _rows(self, read, auth):
+ return read.to_arrow([QueryAuthSplit(object(), auth)],
parallelism=1).to_pydict()
+
+ def test_the_reader_actually_masks(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ read = self._read(fields, {"id": ["1"], "email": ["SECRET"]})
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ self.assertEqual(self._rows(read, auth), {"id": ["1"], "email":
[None]})
+
+ def test_the_reader_applies_the_row_filter(self):
+ fields = [_string_field(0, "id"), _string_field(1, "dept")]
+ read = self._read(fields, {"id": ["1", "2"], "dept": ["eng", "ops"]})
+ auth = TableQueryAuthResult([_filter_json("dept", "eng")], None)
+ self.assertEqual(self._rows(read, auth), {"id": ["1"], "dept":
["eng"]})
+
+ def test_a_mask_on_a_column_added_after_the_snapshot_is_inert(self):
+ fields = [_string_field(0, "id")]
+ latest = [_string_field(0, "id"), _string_field(1, "new_email")]
+ read = self._read(fields, {"id": ["1"]}, latest=latest)
+ auth = TableQueryAuthResult(None, {"new_email": _null_mask()})
+ self.assertEqual(self._rows(read, auth), {"id": ["1"]})
+
+ def test_a_stale_rule_is_refused_at_the_reader(self):
+ fields = [_string_field(0, "id"), _string_field(1, "mail")]
+ read = self._read(fields, {"id": ["1"], "mail": ["SECRET"]})
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ with self.assertRaises(ValueError):
+ self._rows(read, auth)
+
+ def test_a_cross_mask_is_refused_at_the_reader(self):
+ fields = [_string_field(0, "email"), _string_field(1, "alias")]
+ read = self._read(fields, {"email": ["SECRET"], "alias": ["x"]})
+ auth = TableQueryAuthResult(None, {
+ "email": _null_mask(), "alias": _field_ref_mask("email", 0)})
+ with self.assertRaises(ValueError):
+ self._rows(read, auth)
+
+ def
test_an_unprojected_system_filter_operand_is_refused_at_the_reader(self):
+ fields = [_string_field(0, "id")]
+ read = self._read(fields, {"id": ["1"]})
+ auth = TableQueryAuthResult([_filter_json("_ROW_ID")], None)
+ with self.assertRaises(RuntimeError):
+ self._rows(read, auth)
+
+ def
test_a_mask_on_an_unprojected_system_column_is_inert_at_the_reader(self):
+ fields = [_string_field(0, "id")]
+ read = self._read(fields, {"id": ["1"]})
+ auth = TableQueryAuthResult(None, {"_ROW_ID": _null_mask()})
+ self.assertEqual(self._rows(read, auth), {"id": ["1"]})
+
+ def test_a_metadata_read_failure_is_propagated_at_the_reader(self):
+ fields = [_string_field(0, "id"), _string_field(1, "email")]
+ read = self._read(fields, {"id": ["1"], "email": ["SECRET"]})
+
+ def boom():
+ raise RuntimeError("no schema")
+ read.table.schema_manager.latest = boom
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ with self.assertRaises(RuntimeError):
+ self._rows(read, auth)
+
+ def test_a_widened_column_is_read_and_masked(self):
+ historical = [DataField(0, "n", AtomicType("INT"))]
+ latest = [DataField(0, "n", AtomicType("BIGINT"))]
+ read = self._read(historical, {"n": [7]}, latest=latest)
+ auth = TableQueryAuthResult(None, {"n": _null_mask()})
+ self.assertEqual(self._rows(read, auth), {"n": [None]})
+
+ def test_a_generated_alias_colliding_with_a_new_column_is_refused(self):
+ historical = [_string_field(0, "email")]
+ latest = [_string_field(0, "email"), _string_field(3, "email__0")]
+ projected = [_string_field(0, "email"), _string_field(0, "email__0")]
+ read = self._read(historical, {"email": ["SECRET"], "email__0":
["SECRET"]},
+ latest=latest, projected=projected)
+ auth = TableQueryAuthResult(None, {"email__0": _null_mask()})
+ with self.assertRaises((ValueError, RuntimeError)):
+ self._rows(read, auth)
+
+ def test_a_swapped_identity_is_refused_at_the_reader(self):
+ historical = [_string_field(1, "email"), _string_field(2, "alias")]
+ latest = [_string_field(2, "email"), _string_field(1, "alias")]
+ read = self._read(historical, {"email": ["SECRET"], "alias": ["x"]},
+ latest=latest)
+ auth = TableQueryAuthResult(None, {"email": _null_mask()})
+ with self.assertRaises(ValueError):
+ self._rows(read, auth)
diff --git a/paimon-python/pypaimon/tests/read_builder_explain_test.py
b/paimon-python/pypaimon/tests/read_builder_explain_test.py
index e7489d6a49..59d1063861 100644
--- a/paimon-python/pypaimon/tests/read_builder_explain_test.py
+++ b/paimon-python/pypaimon/tests/read_builder_explain_test.py
@@ -21,6 +21,7 @@ split-level execution signals, and pretty-print smoke."""
import os
import shutil
import tempfile
+import json
import unittest
from typing import Any, Dict, List
@@ -255,7 +256,7 @@ class ReadBuilderExplainTest(unittest.TestCase):
table, pa_schema = self._append_table('explain_query_auth')
_write(table, [{'id': i, 'val': i} for i in range(20)], pa_schema)
- auth_result = TableQueryAuthResult(filter=None, column_masking={'val':
'CAST(NULL AS BIGINT)'})
+ auth_result = TableQueryAuthResult(filter=None, column_masking={'val':
json.dumps({'name': 'NULL'})})
table.catalog_environment.table_query_auth = lambda options,
identifier: (lambda select: auth_result)
rb = table.new_read_builder()