paimon-python/pypaimon/tests/table_query_auth_test.py

################################################################################
#  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 unittest
from unittest.mock import patch

import pyarrow as pa
import pytest

from types import SimpleNamespace

from pypaimon import Schema
from pypaimon.api.api_response import AuthTableQueryResponse
from pypaimon.catalog.catalog_exception import TableNoPermissionException
from pypaimon.catalog.table_query_auth import (
    TableQueryAuthResult,
)
from pypaimon.common.identifier import Identifier
from pypaimon.common.options import Options
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.read.plan import Plan
from pypaimon.read.query_auth_split import QueryAuthSplit
from pypaimon.read.scan_stats import ScanStats
from pypaimon.read.table_scan import TableScan
from pypaimon.tests import native_plan_rest_test


rest_catalog = native_plan_rest_test.rest_catalog


class _FakeField:
    def __init__(self, name):
        self.name = name


class _FakeSplit:
    def __init__(self, name="s1"):
        self.name = name

    @property
    def row_count(self):
        return 100

    @property
    def files(self):
        return []

    @property
    def partition(self):
        return {}

    @property
    def bucket(self):
        return 0

    def merged_row_count(self):
        return 100


class _FakeTable:
    def __init__(self, identifier=None):
        self.identifier = identifier or Identifier("db", "table")
        self.options = SimpleNamespace(
            native_plan_enabled=lambda: False, native_read_enabled=lambda: False, query_auth_enabled=False)


class _FakeScan(TableScan):
    def __init__(self, plan, identifier=None, auth_result=None):
        self._plan = plan
        self._query_auth_fn = None
        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, all_columns=False):
        return self._auth_result


class _FakeFileScanner:
    def __init__(self, plan):
        self._plan = plan
        self.skip_level0 = False

    def scan(self):
        return self._plan


def _simple_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 TestTableNoPermissionException(unittest.TestCase):

    def test_message_format(self):
        identifier = Identifier("db", "table")
        exc = TableNoPermissionException(identifier)
        self.assertIn("db.table", str(exc))
        self.assertIn("No permission", str(exc))
        self.assertEqual(exc.identifier, identifier)


class TestTableQueryAuthResultConvertPlan(unittest.TestCase):

    def test_no_auth_returns_original_plan(self):
        result = TableQueryAuthResult(None, None)
        plan = Plan([_FakeSplit()])
        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 = Plan([_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_entry_is_rejected(self):
        # Skipping it would read every row.
        for rules in (["", _simple_filter_json()], [None]):
            with self.assertRaisesRegex(ValueError, "Row filter cannot be empty"):
                TableQueryAuthResult(rules, None)

    def test_blank_column_masking_transform_is_rejected(self):
        # Skipping it would return the column unmasked.
        with self.assertRaisesRegex(ValueError, "Column masking on 'col' cannot be empty"):
            TableQueryAuthResult(None, {"col": "", "col3": '{"name":"NULL"}'})

    def test_blank_column_masking_target_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "Column masking target cannot be empty"):
            TableQueryAuthResult(None, {"": '{"name":"NULL"}'})

    def test_malformed_filter_is_rejected(self):
        for rules in ({}, "", 0, False, {"a": "b"}):
            with self.assertRaisesRegex(ValueError, "Row filter must be a list"):
                TableQueryAuthResult(rules, None)

    def test_malformed_column_masking_is_rejected(self):
        for masking in ([], "", 0, False, ['{"name":"NULL"}']):
            with self.assertRaisesRegex(ValueError, "Column masking must be a map"):
                TableQueryAuthResult(None, masking)

    def test_wraps_splits_with_filter(self):
        result = TableQueryAuthResult([_simple_filter_json()], None)
        stats = ScanStats(entries_after_stats=2)
        plan = Plan([_FakeSplit("s1"), _FakeSplit("s2")], _scan_stats=stats)
        converted = result.convert_plan(plan)
        self.assertIs(converted.stats(), stats)
        self.assertEqual(len(converted.splits()), 2)
        for qs in converted.splits():
            self.assertIsInstance(qs, QueryAuthSplit)

    def test_auth_wrapping_preserves_lazy_native_statistics(self):
        result = TableQueryAuthResult([_simple_filter_json()], None)
        handle = object()
        plan = Plan([_FakeSplit("s1")], snapshot_id=7, _native_stats=handle)
        counters = ScanStats.empty_native()
        counters.entries_total = 10
        counters.planning_counters['manifest_entries_read'] = 10
        with patch.object(ScanStats, 'from_native', return_value=counters) as convert:
            wrapped = result.convert_plan(plan)
            convert.assert_not_called()
            self.assertEqual(wrapped.snapshot_id, 7)
            self.assertIsInstance(wrapped.splits()[0], QueryAuthSplit)
            self.assertIs(wrapped.stats(), counters)
            self.assertIs(wrapped.stats(), counters)
            self.assertEqual(counters.planning_counters['manifest_entries_read'], 10)
            self.assertEqual(counters.planning_counters['final_splits'], 1)
            convert.assert_called_once_with(handle)

    def test_wraps_splits_with_masking(self):
        result = TableQueryAuthResult(None, {"col": '{"name":"NULL"}'})
        plan = Plan([_FakeSplit()])
        converted = result.convert_plan(plan)
        self.assertEqual(len(converted.splits()), 1)
        self.assertIsInstance(converted.splits()[0], QueryAuthSplit)

    def test_inner_split_preserved(self):
        result = TableQueryAuthResult([_simple_filter_json()], None)
        plan = Plan([_FakeSplit("original")])
        converted = result.convert_plan(plan)
        self.assertEqual(converted.splits()[0].split.name, "original")


class TestTableQueryAuthResultExtractRowFilter(unittest.TestCase):

    def test_no_filter_returns_none(self):
        result = TableQueryAuthResult(None, None)
        self.assertIsNone(result.extract_row_filter())

    def test_single_filter(self):
        result = TableQueryAuthResult([_simple_filter_json("dept", "eng")], None)
        fn = result.extract_row_filter()
        self.assertIsNotNone(fn)
        batch = pa.RecordBatch.from_pydict({"dept": ["eng", "sales", "eng"]})
        mask = fn(batch)
        self.assertEqual(mask.to_pylist(), [True, False, True])

    def test_multiple_filters_combined_with_and(self):
        f1 = _simple_filter_json("dept", "eng")
        f2 = json.dumps({
            "kind": "LEAF",
            "transform": {
                "name": "FIELD_REF",
                "fieldRef": {"index": 0, "name": "dept", "type": "STRING"},
            },
            "function": "NOT_EQUAL",
            "literals": ["eng"],
        })
        result = TableQueryAuthResult([f1, f2], None)
        fn = result.extract_row_filter()
        batch = pa.RecordBatch.from_pydict({"dept": ["eng", "sales", "eng"]})
        mask = fn(batch)
        self.assertEqual(mask.to_pylist(), [False, False, False])

    def test_filter_entries_short_circuit_in_server_order(self):
        first = _simple_filter_json('dept', 'eng')
        second = json.dumps({
            'kind': 'LEAF', 'function': 'EQUAL', 'literals': ['a'],
            'transform': {'name': 'SUBSTRING', 'inputs': [
                {'name': 's', 'type': 'STRING'}, {'name': 'position', 'type': 'BIGINT'}, 1]},
        })
        batch = pa.RecordBatch.from_pydict({
            'dept': ['eng', 'sales'], 's': ['abc', 'xyz'], 'position': [1, 2147483648]})
        result = TableQueryAuthResult([first, second], None).extract_row_filter()
        self.assertEqual(result(batch).to_pylist(), [True, False])
        # Reversing the response's order executes the throwing transform first.
        with self.assertRaises(ValueError):
            TableQueryAuthResult([second, first], None).extract_row_filter()(batch)


class TestTableQueryAuthResultExtraFields(unittest.TestCase):

    def test_detects_unprojected_field(self):
        read_fields = [_FakeField("name"), _FakeField("age")]
        table_fields = [
            _FakeField("name"),
            _FakeField("age"),
            _FakeField("dept"),
        ]
        result = TableQueryAuthResult([_simple_filter_json("dept")], None)
        extra = result.get_extra_fields(read_fields, table_fields)
        self.assertEqual(len(extra), 1)
        self.assertEqual(extra[0].name, "dept")

    def test_no_extra_when_already_projected(self):
        read_fields = [_FakeField("name"), _FakeField("dept")]
        table_fields = read_fields + [_FakeField("age")]
        result = TableQueryAuthResult([_simple_filter_json("dept")], None)
        extra = result.get_extra_fields(read_fields, table_fields)
        self.assertEqual(len(extra), 0)

    def test_no_extra_when_no_filter(self):
        result = TableQueryAuthResult(None, None)
        extra = result.get_extra_fields(
            [_FakeField("a")], [_FakeField("a"), _FakeField("b")]
        )
        self.assertEqual(len(extra), 0)

    def test_deduplicates_extra_fields(self):
        f1 = _simple_filter_json("dept", "eng")
        f2 = _simple_filter_json("dept", "sales")
        read_fields = [_FakeField("name")]
        table_fields = [_FakeField("name"), _FakeField("dept")]
        result = TableQueryAuthResult([f1, f2], None)
        extra = result.get_extra_fields(read_fields, table_fields)
        self.assertEqual(len(extra), 1)


class TestQueryAuthSplit(unittest.TestCase):

    def test_delegates_properties(self):
        auth = TableQueryAuthResult([_simple_filter_json()], None)
        split = _FakeSplit()
        qs = QueryAuthSplit(split, auth)
        self.assertEqual(qs.row_count, 100)
        self.assertEqual(qs.bucket, 0)
        self.assertEqual(qs.files, [])
        self.assertEqual(qs.partition, {})

    def test_merged_row_count_none_with_filter(self):
        auth = TableQueryAuthResult([_simple_filter_json()], None)
        qs = QueryAuthSplit(_FakeSplit(), auth)
        self.assertIsNone(qs.merged_row_count())

    def test_merged_row_count_delegates_without_filter(self):
        auth = TableQueryAuthResult(None, {"col": '{"name":"NULL"}'})
        qs = QueryAuthSplit(_FakeSplit(), auth)
        self.assertEqual(qs.merged_row_count(), 100)

    def test_exposes_inner_split(self):
        split = _FakeSplit("inner")
        qs = QueryAuthSplit(split, TableQueryAuthResult(None, None))
        self.assertIs(qs.split, split)

    def test_exposes_auth_result(self):
        auth = TableQueryAuthResult(None, None)
        qs = QueryAuthSplit(_FakeSplit(), auth)
        self.assertIs(qs.auth_result, auth)


class TestQueryAuthSplitRawConvertible(unittest.TestCase):

    def test_raw_convertible_proxy_true(self):
        inner = _FakeSplit()
        inner.raw_convertible = True
        auth_split = QueryAuthSplit(inner, None)
        self.assertTrue(auth_split.raw_convertible)

    def test_raw_convertible_proxy_false(self):
        inner = _FakeSplit()
        inner.raw_convertible = False
        auth_split = QueryAuthSplit(inner, None)
        self.assertFalse(auth_split.raw_convertible)


class TestQueryAuthSplitAttributeDelegation(unittest.TestCase):

    def test_file_size_delegated(self):
        inner = _FakeSplit()
        inner.file_size = 12345
        auth_split = QueryAuthSplit(inner, TableQueryAuthResult(None, None))
        self.assertEqual(auth_split.file_size, 12345)

    def test_file_paths_delegated(self):
        inner = _FakeSplit()
        inner.file_paths = ["/data/file1.parquet", "/data/file2.parquet"]
        auth_split = QueryAuthSplit(inner, TableQueryAuthResult(None, None))
        self.assertEqual(auth_split.file_paths, ["/data/file1.parquet", "/data/file2.parquet"])

    def test_data_deletion_files_delegated(self):
        inner = _FakeSplit()
        inner.data_deletion_files = [None, "dv-1"]
        auth_split = QueryAuthSplit(inner, TableQueryAuthResult(None, None))
        self.assertEqual(auth_split.data_deletion_files, [None, "dv-1"])

    def test_unknown_attr_raises(self):
        inner = _FakeSplit()
        auth_split = QueryAuthSplit(inner, TableQueryAuthResult(None, None))
        with self.assertRaises(AttributeError):
            _ = auth_split.completely_nonexistent_attr

    def test_pickle_roundtrip(self):
        import pickle
        inner = _FakeSplit()
        inner.file_size = 999
        auth_split = QueryAuthSplit(inner, TableQueryAuthResult(None, None))
        restored = pickle.loads(pickle.dumps(auth_split))
        self.assertEqual(restored.file_size, 999)
        self.assertEqual(restored.row_count, 100)


class TestPlanAllFiles(unittest.TestCase):

    def test_no_restriction_returns_plan_unchanged(self):
        from pypaimon.read.table_scan import TableScan

        plain_plan = Plan([_FakeSplit()], snapshot_id=5)
        scan = _FakeScan(plain_plan, auth_result=None)
        result = TableScan.plan_all(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)
        plan = Plan([_FakeSplit()], snapshot_id=5)
        scan = _FakeScan(plan, auth_result=auth)
        with self.assertRaises(TableNoPermissionException):
            TableScan.plan_all(scan)

    def test_column_masking_restriction_raises(self):
        from pypaimon.read.table_scan import TableScan

        auth = TableQueryAuthResult(None, {"col": '{"name":"NULL"}'})
        plan = Plan([_FakeSplit()], snapshot_id=5)
        scan = _FakeScan(plan, auth_result=auth)
        with self.assertRaises(TableNoPermissionException):
            TableScan.plan_all(scan)

    def test_empty_table_with_auth_still_raises(self):
        from pypaimon.read.table_scan import TableScan

        auth = TableQueryAuthResult([_simple_filter_json()], None)
        empty_plan = Plan([], snapshot_id=5)
        scan = _FakeScan(empty_plan, auth_result=auth)
        with self.assertRaises(TableNoPermissionException):
            TableScan.plan_all(scan)

    def test_no_restrictions_allows_write(self):
        from pypaimon.read.table_scan import TableScan

        plan = Plan([_FakeSplit()], snapshot_id=5)
        scan = _FakeScan(plan, auth_result=None)
        result = TableScan.plan_all(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_fails_resolution(self):
        from pypaimon.read.query_auth_split import resolve_auth_result
        fn = lambda select: TableQueryAuthResult(["", None], None)
        with self.assertRaisesRegex(ValueError, "Row filter cannot be empty"):
            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):

    def test_disabled_by_default(self):
        opts = CoreOptions(Options({}))
        self.assertFalse(opts.query_auth_enabled)

    def test_enabled_when_set(self):
        opts = CoreOptions(Options({"query-auth.enabled": True}))
        self.assertTrue(opts.query_auth_enabled)

    def test_disabled_when_false(self):
        opts = CoreOptions(Options({"query-auth.enabled": False}))
        self.assertFalse(opts.query_auth_enabled)


@pytest.mark.parametrize("copy_method", ["copy", "copy_without_time_travel"])
@pytest.mark.parametrize("enabled", ["true", True, "on"])
def test_query_auth_cannot_be_disabled_dynamically(rest_catalog, copy_method, enabled):
    catalog, server = rest_catalog
    schema = pa.schema([("dept", pa.string()), ("secret", pa.string())])
    catalog.create_table("default.auth", Schema.from_pyarrow_schema(
        schema, options={"query-auth.enabled": enabled}), False)
    table = catalog.get_table("default.auth")
    builder = table.new_batch_write_builder()
    writer = builder.new_write()
    commit = builder.new_commit()
    try:
        writer.write_arrow(pa.Table.from_pydict({
            "dept": ["eng", "other"], "secret": ["secret1", "secret2"],
        }, schema=schema))
        commit.commit(writer.prepare_commit())
    finally:
        writer.close()
        commit.close()
    server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse(
        filter=[_simple_filter_json()], column_masking={"secret": '{"name":"NULL"}'})

    copy = getattr(table, copy_method)
    for value in ("false", "FaLsE", False, None, "", "invalid", "0", 0, "off", "no"):
        with pytest.raises(ValueError) as error:
            copy({"query-auth.enabled": value})
        assert str(error.value) == "Option 'query-auth.enabled' is not allowed to be disabled dynamically"

    for options in [{}] + [{"query-auth.enabled": value} for value in ("true", "TrUe", True, " on ", 1)]:
        copied = copy(options)
        read = copied.new_read_builder()
        result = read.new_read().to_arrow(read.new_scan().plan().splits())
        assert result.to_pydict() == {"dept": ["eng"], "secret": [None]}
        with pytest.raises(ValueError, match="Search is not supported on a query-auth table"):
            copied.new_vector_search_builder().new_vector_search_scan()
    assert table.table_schema.options["query-auth.enabled"] == enabled


@pytest.mark.parametrize("copy_method", ["copy", "copy_without_time_travel"])
@pytest.mark.parametrize("options", [{}, {"query-auth.enabled": "false"}])
def test_query_auth_can_be_enabled_dynamically(rest_catalog, copy_method, options):
    catalog, _ = rest_catalog
    catalog.create_table("default.plain", Schema.from_pyarrow_schema(
        pa.schema([("secret", pa.string())]), options=options), False)
    table = catalog.get_table("default.plain")
    for value, enabled in (("false", False), (False, False), (None, False),
                           ("true", True), ("TrUe", True), (True, True), ("on", True)):
        copied = getattr(table, copy_method)({"query-auth.enabled": value})
        assert copied.options.query_auth_enabled == enabled
    assert not table.options.query_auth_enabled


if __name__ == "__main__":
    unittest.main()


class TestConvertPlanPreservesSnapshotId(unittest.TestCase):

    def test_convert_plan_preserves_snapshot_id(self):
        from pypaimon.catalog.table_query_auth import TableQueryAuthResult
        from pypaimon.read.plan import Plan
        from unittest.mock import MagicMock

        split = MagicMock()
        original_plan = Plan([split], snapshot_id=42)
        predicate_json = (
            '{"kind":"LEAF","transform":{"name":"FIELD_REF",'
            '"fieldRef":{"name":"id"}},"function":"EQUAL","literals":[1]}'
        )
        auth_result = TableQueryAuthResult(
            filter=[predicate_json],
            column_masking=None)
        converted = auth_result.convert_plan(original_plan)
        assert converted.snapshot_id == 42

    def test_convert_plan_preserves_none_snapshot_id(self):
        from pypaimon.catalog.table_query_auth import TableQueryAuthResult
        from pypaimon.read.plan import Plan
        from unittest.mock import MagicMock

        split = MagicMock()
        original_plan = Plan([split], snapshot_id=None)
        predicate_json = (
            '{"kind":"LEAF","transform":{"name":"FIELD_REF",'
            '"fieldRef":{"name":"id"}},"function":"EQUAL","literals":[1]}'
        )
        auth_result = TableQueryAuthResult(
            filter=[predicate_json],
            column_masking=None)
        converted = auth_result.convert_plan(original_plan)
        assert converted.snapshot_id is None