[python] Fix write-path auth bypass and unify auth restriction check

Author: MgjLLL

paimon-python/pypaimon/catalog/table_query_auth.py

@@ -25,6 +25,8 @@
     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
 
 
@@ -37,11 +39,12 @@ def __init__(self, filter: Optional[List[str]], column_masking: Optional[Dict[st
             if column_masking else column_masking
         )
 
-    def convert_plan(self, plan):
-        from pypaimon.read.query_auth_split import QueryAuthSplit
-        from pypaimon.read.plan import Plan
+    @property
+    def has_restrictions(self):
+        return bool(self.filter) or bool(self.column_masking)
 
-        if not self.filter and not self.column_masking:
+    def convert_plan(self, plan):
+        if not self.has_restrictions:
             return plan
         auth_splits = [QueryAuthSplit(split, self) for split in plan.splits()]
         return Plan(auth_splits, snapshot_id=plan.snapshot_id)

paimon-python/pypaimon/read/query_auth_split.py

@@ -68,7 +68,10 @@ def resolve_auth_result(query_auth_fn, read_type):
     if query_auth_fn is None:
         return None
     select = [f.name for f in read_type] if read_type else None
-    return query_auth_fn(select)
+    result = query_auth_fn(select)
+    if result is None or not result.has_restrictions:
+        return None
+    return result
 
 
 def wrap_plan_with_auth(auth_result, plan):

paimon-python/pypaimon/read/table_scan.py

@@ -24,7 +24,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 QueryAuthSplit, resolve_auth_result, wrap_plan_with_auth
+from pypaimon.read.query_auth_split import resolve_auth_result, wrap_plan_with_auth
 from pypaimon.read.scan_stats import ScanStats
 from pypaimon.read.scanner.file_scanner import FileScanner
 
@@ -58,10 +58,9 @@ def plan(self) -> Plan:
         return wrap_plan_with_auth(auth_result, plan)
 
     def plan_for_write(self) -> Plan:
-        plan = self.plan()
-        if any(isinstance(s, QueryAuthSplit) for s in plan.splits()):
+        if self.__auth_query() is not None:
             raise TableNoPermissionException(self.table.identifier)
-        return plan
+        return self.file_scanner.scan()
 
     def __auth_query(self):
         return resolve_auth_result(self._query_auth_fn, self._read_type)

paimon-python/pypaimon/tests/table_query_auth_test.py

@@ -75,13 +75,26 @@ def __init__(self, identifier=None):
 
 
 class _FakeScan:
-    def __init__(self, plan, identifier=None):
+    def __init__(self, plan, identifier=None, auth_result=None):
         self._plan = plan
         self.table = _FakeTable(identifier)
+        self.file_scanner = _FakeFileScanner(plan)
+        self._auth_result = auth_result
 
     def plan(self):
         return self._plan
 
+    def _TableScan__auth_query(self):
+        return self._auth_result
+
+
+class _FakeFileScanner:
+    def __init__(self, plan):
+        self._plan = plan
+
+    def scan(self):
+        return self._plan
+
 
 def _simple_filter_json(field_name="dept", value="eng"):
     return json.dumps({
@@ -113,12 +126,28 @@ def test_no_auth_returns_original_plan(self):
         converted = result.convert_plan(plan)
         self.assertIs(converted, plan)
 
+    def test_no_auth_has_no_restrictions(self):
+        result = TableQueryAuthResult(None, None)
+        self.assertFalse(result.has_restrictions)
+
     def test_empty_filter_and_masking_returns_original(self):
         result = TableQueryAuthResult([], {})
         plan = _FakePlan([_FakeSplit()])
         converted = result.convert_plan(plan)
         self.assertIs(converted, plan)
 
+    def test_empty_filter_and_masking_has_no_restrictions(self):
+        result = TableQueryAuthResult([], {})
+        self.assertFalse(result.has_restrictions)
+
+    def test_filter_has_restrictions(self):
+        result = TableQueryAuthResult([_simple_filter_json()], None)
+        self.assertTrue(result.has_restrictions)
+
+    def test_masking_has_restrictions(self):
+        result = TableQueryAuthResult(None, {"col": '{"name":"NULL"}'})
+        self.assertTrue(result.has_restrictions)
+
     def test_blank_filter_entries_are_skipped(self):
         result = TableQueryAuthResult(["", None], None)
         self.assertFalse(result.filter)
@@ -334,38 +363,81 @@ def test_no_restriction_returns_plan_unchanged(self):
         from pypaimon.read.table_scan import TableScan
 
         plain_plan = _FakePlan([_FakeSplit()], snapshot_id=5)
-        scan = _FakeScan(plain_plan)
+        scan = _FakeScan(plain_plan, auth_result=None)
         result = TableScan.plan_for_write(scan)
         self.assertIs(result, plain_plan)
 
     def test_row_filter_restriction_raises(self):
         from pypaimon.read.table_scan import TableScan
 
         auth = TableQueryAuthResult([_simple_filter_json()], None)
-        wrapped_plan = _FakePlan([QueryAuthSplit(_FakeSplit(), auth)], snapshot_id=5)
-        scan = _FakeScan(wrapped_plan)
+        plan = _FakePlan([_FakeSplit()], snapshot_id=5)
+        scan = _FakeScan(plan, auth_result=auth)
         with self.assertRaises(TableNoPermissionException):
             TableScan.plan_for_write(scan)
 
     def test_column_masking_restriction_raises(self):
         from pypaimon.read.table_scan import TableScan
 
         auth = TableQueryAuthResult(None, {"col": '{"name":"NULL"}'})
-        wrapped_plan = _FakePlan([QueryAuthSplit(_FakeSplit(), auth)], snapshot_id=5)
-        scan = _FakeScan(wrapped_plan)
+        plan = _FakePlan([_FakeSplit()], snapshot_id=5)
+        scan = _FakeScan(plan, auth_result=auth)
         with self.assertRaises(TableNoPermissionException):
             TableScan.plan_for_write(scan)
 
-    def test_mixed_splits_any_restricted_raises(self):
+    def test_empty_table_with_auth_still_raises(self):
         from pypaimon.read.table_scan import TableScan
 
         auth = TableQueryAuthResult([_simple_filter_json()], None)
-        wrapped_plan = _FakePlan(
-            [_FakeSplit(), QueryAuthSplit(_FakeSplit(), auth)], snapshot_id=5)
-        scan = _FakeScan(wrapped_plan)
+        empty_plan = _FakePlan([], snapshot_id=5)
+        scan = _FakeScan(empty_plan, auth_result=auth)
         with self.assertRaises(TableNoPermissionException):
             TableScan.plan_for_write(scan)
 
+    def test_no_restrictions_allows_write(self):
+        from pypaimon.read.table_scan import TableScan
+
+        plan = _FakePlan([_FakeSplit()], snapshot_id=5)
+        scan = _FakeScan(plan, auth_result=None)
+        result = TableScan.plan_for_write(scan)
+        self.assertIs(result, plan)
+
+
+class TestResolveAuthResult(unittest.TestCase):
+
+    def test_none_fn_returns_none(self):
+        from pypaimon.read.query_auth_split import resolve_auth_result
+        self.assertIsNone(resolve_auth_result(None, None))
+
+    def test_no_restrictions_returns_none(self):
+        from pypaimon.read.query_auth_split import resolve_auth_result
+        fn = lambda select: TableQueryAuthResult(None, None)
+        self.assertIsNone(resolve_auth_result(fn, None))
+
+    def test_empty_filter_returns_none(self):
+        from pypaimon.read.query_auth_split import resolve_auth_result
+        fn = lambda select: TableQueryAuthResult([], {})
+        self.assertIsNone(resolve_auth_result(fn, None))
+
+    def test_blank_filter_stripped_returns_none(self):
+        from pypaimon.read.query_auth_split import resolve_auth_result
+        fn = lambda select: TableQueryAuthResult(["", None], None)
+        self.assertIsNone(resolve_auth_result(fn, None))
+
+    def test_with_filter_returns_result(self):
+        from pypaimon.read.query_auth_split import resolve_auth_result
+        fn = lambda select: TableQueryAuthResult([_simple_filter_json()], None)
+        result = resolve_auth_result(fn, None)
+        self.assertIsNotNone(result)
+        self.assertTrue(result.has_restrictions)
+
+    def test_with_masking_returns_result(self):
+        from pypaimon.read.query_auth_split import resolve_auth_result
+        fn = lambda select: TableQueryAuthResult(None, {"col": '{"name":"NULL"}'})
+        result = resolve_auth_result(fn, None)
+        self.assertIsNotNone(result)
+        self.assertTrue(result.has_restrictions)
+
 
 class TestCoreOptionsQueryAuth(unittest.TestCase):