diff --git a/paimon-python/pypaimon/catalog/table_query_auth.py b/paimon-python/pypaimon/catalog/table_query_auth.py index 5a546d284868..523c42742df6 100644 --- a/paimon-python/pypaimon/catalog/table_query_auth.py +++ b/paimon-python/pypaimon/catalog/table_query_auth.py @@ -30,6 +30,17 @@ from pypaimon.schema.data_types import DataField +def reject_search_under_query_auth(table) -> None: + """Refuses a search on a query-auth table. Called from the methods that build a scan or a + read rather than from the builder constructors, which a deserialized builder skips.""" + from pypaimon.table.file_store_table import FileStoreTable + + if isinstance(table, FileStoreTable) and table.options.query_auth_enabled: + raise ValueError( + "Search is not supported on a query-auth table: the index ranks raw values, " + "which a column mask invalidates.") + + class TableQueryAuthResult: def __init__(self, filter: Optional[List[str]], column_masking: Optional[Dict[str, str]]): diff --git a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py index ce2f3ed9d4f4..b07cd1e835cd 100644 --- a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py +++ b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py @@ -19,6 +19,7 @@ from abc import ABC, abstractmethod +from pypaimon.catalog.table_query_auth import reject_search_under_query_auth from pypaimon.table.source.vector_search_builder import ( AbstractVectorSearchBuilderImpl, ) @@ -110,6 +111,7 @@ def with_query_vectors(self, vectors): def new_batch_vector_search_read(self): # type: () -> BatchVectorSearchRead + reject_search_under_query_auth(self._table) if self._limit <= 0: raise ValueError("Limit must be positive, set via with_limit()") if self._vector_column is None: diff --git a/paimon-python/pypaimon/table/source/full_text_search_builder.py b/paimon-python/pypaimon/table/source/full_text_search_builder.py index 1e094e543f62..420f918ac5cf 100644 --- a/paimon-python/pypaimon/table/source/full_text_search_builder.py +++ b/paimon-python/pypaimon/table/source/full_text_search_builder.py @@ -20,6 +20,7 @@ from abc import ABC, abstractmethod from typing import Optional +from pypaimon.catalog.table_query_auth import reject_search_under_query_auth from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.globalindex.global_index_result import GlobalIndexResult from pypaimon.table.source.full_text_read import FullTextRead, DataEvolutionFullTextRead @@ -118,6 +119,7 @@ def _rebuild_leaf_indices_by_name(cls, predicate, name_to_idx): return predicate.new_index(name_to_idx[predicate.field]) def new_full_text_scan(self) -> FullTextScan: + reject_search_under_query_auth(self._table) definition = self._primary_key_full_text_definition() if definition is not None: from pypaimon.table.source.primary_key_full_text_scan import PrimaryKeyFullTextScan @@ -130,6 +132,7 @@ def new_full_text_scan(self) -> FullTextScan: ) def new_full_text_read(self) -> FullTextRead: + reject_search_under_query_auth(self._table) if self._limit <= 0: raise ValueError("Limit must be positive, set via with_limit()") definition = self._primary_key_full_text_definition() diff --git a/paimon-python/pypaimon/table/source/hybrid_search_builder.py b/paimon-python/pypaimon/table/source/hybrid_search_builder.py index a3e43c3fb6cc..01736816dbfa 100644 --- a/paimon-python/pypaimon/table/source/hybrid_search_builder.py +++ b/paimon-python/pypaimon/table/source/hybrid_search_builder.py @@ -25,6 +25,7 @@ from dataclasses import dataclass, field from typing import Dict, List, Optional +from pypaimon.catalog.table_query_auth import reject_search_under_query_auth from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.globalindex.global_index_result import GlobalIndexResult from pypaimon.globalindex.vector_search_result import ( @@ -312,6 +313,7 @@ def add_route(self, route: HybridSearchRoute) -> 'HybridSearchBuilder': return self def route_builders(self) -> List[HybridSearchRouteBuilder]: + reject_search_under_query_auth(self._table) self._validate_search() from pypaimon.snapshot.time_travel_util import TimeTravelUtil execution = copy(self) diff --git a/paimon-python/pypaimon/table/source/vector_search_builder.py b/paimon-python/pypaimon/table/source/vector_search_builder.py index af823e58c2fb..2c7e309a3f47 100644 --- a/paimon-python/pypaimon/table/source/vector_search_builder.py +++ b/paimon-python/pypaimon/table/source/vector_search_builder.py @@ -19,6 +19,7 @@ from abc import ABC, abstractmethod +from pypaimon.catalog.table_query_auth import reject_search_under_query_auth from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.table.source.vector_search_read import DataEvolutionVectorRead from pypaimon.table.source.vector_search_scan import DataEvolutionVectorScan @@ -215,6 +216,7 @@ def _rebuild_leaf_indices_by_name(cls, predicate, pk_to_idx): def new_vector_search_scan(self): # type: () -> VectorSearchScan + reject_search_under_query_auth(self._table) if self._vector_column is None: raise ValueError("Vector column must be set via with_vector_column()") scan_class = DataEvolutionVectorScan @@ -250,6 +252,7 @@ def with_query_vector(self, vector): def new_vector_search_read(self): # type: () -> VectorSearchRead + reject_search_under_query_auth(self._table) if self._limit <= 0: raise ValueError("Limit must be positive, set via with_limit()") if self._vector_column is None: diff --git a/paimon-python/pypaimon/tests/search_query_auth_reject_test.py b/paimon-python/pypaimon/tests/search_query_auth_reject_test.py new file mode 100644 index 000000000000..ddc3cea51538 --- /dev/null +++ b/paimon-python/pypaimon/tests/search_query_auth_reject_test.py @@ -0,0 +1,96 @@ +################################################################################ +# 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 os +import shutil +import tempfile +import unittest + +import pyarrow as pa + +from pypaimon import CatalogFactory, Schema +from pypaimon.table.source.batch_vector_search_builder import ( + BatchVectorSearchBuilderImpl, +) +from pypaimon.table.source.full_text_search_builder import FullTextSearchBuilderImpl +from pypaimon.table.source.hybrid_search_builder import HybridSearchBuilderImpl +from pypaimon.table.source.vector_search_builder import VectorSearchBuilderImpl + +REJECTED = "Search is not supported on a query-auth table" + + +def _entry_points(table): + return { + "vector scan": VectorSearchBuilderImpl(table).new_vector_search_scan, + "vector read": VectorSearchBuilderImpl(table).new_vector_search_read, + "batch vector scan": BatchVectorSearchBuilderImpl( + table).new_vector_search_scan, + "batch vector read": BatchVectorSearchBuilderImpl( + table).new_batch_vector_search_read, + "full-text scan": FullTextSearchBuilderImpl(table).new_full_text_scan, + "full-text read": FullTextSearchBuilderImpl(table).new_full_text_read, + "hybrid routes": HybridSearchBuilderImpl(table).route_builders, + } + + +class TestSearchRejectedUnderQueryAuth(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.tempdir = tempfile.mkdtemp() + cls.catalog = CatalogFactory.create( + {'warehouse': os.path.join(cls.tempdir, 'warehouse')}) + cls.catalog.create_database('db', False) + pa_schema = pa.schema([ + ('id', pa.int32()), + ('embedding', pa.list_(pa.float32())), + ('text', pa.string()), + ]) + cls.catalog.create_table( + 'db.plain', Schema.from_pyarrow_schema(pa_schema), False) + cls.catalog.create_table( + 'db.authed', + Schema.from_pyarrow_schema( + pa_schema, options={'query-auth.enabled': 'true'}), + False) + cls.plain = cls.catalog.get_table('db.plain') + cls.authed = cls.catalog.get_table('db.authed') + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.tempdir, ignore_errors=True) + + def test_every_scan_and_read_entry_point_is_rejected(self): + self.assertTrue(self.authed.options.query_auth_enabled) + for name, entry in _entry_points(self.authed).items(): + with self.subTest(entry=name): + with self.assertRaises(ValueError) as ctx: + entry() + self.assertIn(REJECTED, str(ctx.exception)) + + def test_without_query_auth_the_builders_reach_their_own_validation(self): + self.assertFalse(self.plain.options.query_auth_enabled) + for name, entry in _entry_points(self.plain).items(): + with self.subTest(entry=name): + with self.assertRaises(ValueError) as ctx: + entry() + self.assertNotIn(REJECTED, str(ctx.exception)) + + +if __name__ == '__main__': + unittest.main()