[python] Test permissions and query auth across PyPaimon and DataFusion SQL

The mock REST server gains the permission, policy and AuthTableQuery
endpoints, after Java RESTCatalogServer:
- Grants are removed or moved with the table, database or function they
  name, and COLUMN grants follow column renames.
- Per-table policies need query-auth.enabled, conflict as POLICY, must
  name existing columns and go away with their table.
- A schema change is refused before it is committed if it would disable
  query authorization or drop or rename a column that a policy or the
  last allowed column of a COLUMN grant needs.
- The auth endpoint checks the selected columns against the COLUMN grants
  of the principals the caller acts as, and answers with their policies;
  two masks on one column deny the query.

sql_permission_test.py runs in the Rust main job and checks that the two
clients agree: grants and policies made with the CALL sys.* procedures are
listed by PyPaimon and the other way round, and both readers return the
same rows and values under a row filter, column masks and a COLUMN grant.
rest_permission_server_test.py covers the mock server's rules.

Author: plusplusjiajia

paimon-python/pypaimon/tests/rest/rest_permission_server_test.py

@@ -0,0 +1,197 @@
+# 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 pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.api.api_response import ConfigResponse
+from pypaimon.api.auth import BearTokenAuthProvider
+from pypaimon.api.rest_exception import (AlreadyExistsException, BadRequestException,
+                                         NoSuchResourceException)
+from pypaimon.catalog.catalog_exception import TableNoPermissionException
+from pypaimon.common.identifier import Identifier
+from pypaimon.management.column_mask import ColumnMask
+from pypaimon.management.data_policy import DataPolicy
+from pypaimon.management.list_permissions_request import ListPermissionsRequest
+from pypaimon.management.list_policies_request import ListPoliciesRequest
+from pypaimon.management.permission_assignment import PermissionAssignment
+from pypaimon.management.permission_columns import PermissionColumns
+from pypaimon.management.permission_resource import PermissionResource
+from pypaimon.management.resource_type import ResourceType
+from pypaimon.schema.schema_change import SchemaChange
+from pypaimon.tests.rest.rest_function_test import _mock_function
+from pypaimon.tests.rest.rest_server import RESTCatalogServer
+
+NULL_MASK = json.dumps({'name': 'NULL'})
+STARS_MASK = json.dumps({'name': 'CONCAT', 'inputs': ['***']})
+USERS = Identifier.create('default', 'users')
+COLUMNS = PermissionResource(ResourceType.COLUMN, 'default', 'users')
+
+
+def _schema(query_auth=True):
+    return Schema.from_pyarrow_schema(pa.schema([('id', pa.int64()), ('email', pa.string())]),
+                                      options={'query-auth.enabled': 'true'} if query_auth else {})
+
+
+@pytest.fixture
+def rest(tmp_path):
+    server = RESTCatalogServer(str(tmp_path), BearTokenAuthProvider('token'),
+                               ConfigResponse(defaults={'prefix': 'mock'}), 'warehouse')
+    server.start()
+    try:
+        catalog = CatalogFactory.create({
+            'metastore': 'rest', 'uri': server.get_url(), 'warehouse': 'warehouse',
+            'token.provider': 'bear', 'token': 'token', 'data-token.enabled': 'false'})
+        catalog.create_database('default', False)
+        catalog.create_table('default.users', _schema(), False)
+        yield catalog, server
+    finally:
+        server.shutdown()
+
+
+def _table(name):
+    return PermissionResource(ResourceType.TABLE, 'default', name)
+
+
+def _restrict(catalog, name):
+    catalog.permission_management().grant_permission(PermissionAssignment(_table(name), 'SELECT', 'analyst'))
+    catalog.policy_management().create_policy(
+        DataPolicy.column_mask(_table(name), ColumnMask('email', NULL_MASK), 'analyst'))
+
+
+def _access(catalog, name):
+    grants = catalog.permission_management().list_permissions(
+        ListPermissionsRequest(ResourceType.TABLE, 'default', name)).elements
+    policies = catalog.policy_management().list_policies(ListPoliciesRequest(_table(name))).elements
+    return ([(grant.get_access(), grant.get_principal()) for grant in grants],
+            [(policy.get_resource().get_table(), policy.get_column_mask().get_on_column())
+             for policy in policies])
+
+
+def test_table_grants_and_policies_follow_rename_and_drop(rest):
+    catalog, _ = rest
+    _restrict(catalog, 'users')
+
+    catalog.rename_table('default.users', 'default.members')
+    assert _access(catalog, 'members') == ([('SELECT', 'analyst')], [('members', 'email')])
+    catalog.create_table('default.users', _schema(), False)
+    assert _access(catalog, 'users') == ([], [])
+
+    catalog.drop_table('default.members')
+    catalog.create_table('default.members', _schema(), False)
+    assert _access(catalog, 'members') == ([], [])
+
+
+def test_dropping_a_database_or_function_removes_its_grants(rest):
+    catalog, server = rest
+    area = Identifier.create('default', 'area')
+    catalog.create_function(area, _mock_function(area), False)
+    catalog.create_database('sales', False)
+    for resource, access in ((PermissionResource(ResourceType.DATABASE, 'sales'), 'DESCRIBE'),
+                             (PermissionResource(ResourceType.FUNCTION, 'default', function='area'), 'SELECT'),
+                             (PermissionResource(ResourceType.DATABASE, 'default'), 'DESCRIBE')):
+        catalog.permission_management().grant_permission(PermissionAssignment(resource, access, 'analyst'))
+
+    catalog.drop_database('sales')
+    catalog.drop_function(area, False)
+    assert [(grant.get_resource().get_type(), grant.get_resource().get_database())
+            for grant in server.permission_store.values()] == [(ResourceType.DATABASE, 'default')]
+
+
+@pytest.mark.parametrize('expire_time, expired', [
+    (None, False),
+    ('2099-01-01T00:00:00.123000000Z', False),
+    ('2099-01-01t00:00:00.5z', False),
+    ('+12000-01-01T00:00:00Z', False),
+    ('2020-01-01T24:00:00.000000000Z', True),
+    ('-0001-01-01T00:00:00.999Z', True),
+])
+def test_column_grants_limit_the_selected_columns_until_they_expire(rest, expire_time, expired):
+    catalog, server = rest
+    catalog.permission_management().grant_permission(PermissionAssignment(
+        COLUMNS, 'SELECT', 'analyst', PermissionColumns(excluded_column_names=['email']), expire_time))
+
+    server.set_query_principals('analyst')
+    catalog.auth_table_query(USERS, ['id'])
+    for select in (['email'], None):
+        if expired:
+            catalog.auth_table_query(USERS, select)
+        else:
+            with pytest.raises(TableNoPermissionException):
+                catalog.auth_table_query(USERS, select)
+
+
+def test_column_grants_and_policies_need_query_auth_and_existing_columns(rest):
+    catalog, _ = rest
+    catalog.create_table('default.plain', _schema(query_auth=False), False)
+    with pytest.raises(AlreadyExistsException):
+        catalog.permission_management().grant_permission(PermissionAssignment(
+            PermissionResource(ResourceType.COLUMN, 'default', 'plain'), 'SELECT', 'analyst',
+            PermissionColumns(column_names=['id'])))
+    with pytest.raises(NoSuchResourceException):
+        catalog.permission_management().grant_permission(PermissionAssignment(
+            COLUMNS, 'SELECT', 'analyst', PermissionColumns(column_names=['phone'])))
+    with pytest.raises(BadRequestException):
+        catalog.policy_management().create_policy(
+            DataPolicy.column_mask(_table('users'), ColumnMask('phone', NULL_MASK), 'analyst'))
+
+
+def test_schema_changes_cannot_orphan_policies(rest):
+    catalog, _ = rest
+    _restrict(catalog, 'users')
+    schema = catalog.get_table('default.users').table_schema
+    for change in (SchemaChange.set_option('query-auth.enabled', 'false'),
+                   SchemaChange.remove_option('query-auth.enabled'),
+                   SchemaChange.rename_column('email', 'contact'),
+                   SchemaChange.drop_column('email')):
+        with pytest.raises(BadRequestException):
+            catalog.alter_table('default.users', [change])
+
+    current = catalog.get_table('default.users').table_schema
+    assert (current.id, current.options) == (schema.id, schema.options)
+    assert _access(catalog, 'users')[1] == [('users', 'email')]
+    catalog.alter_table('default.users', [SchemaChange.rename_column('id', 'user_id')])
+
+
+def test_column_grants_follow_column_renames(rest):
+    catalog, _ = rest
+    catalog.permission_management().grant_permission(PermissionAssignment(
+        COLUMNS, 'SELECT', 'analyst', PermissionColumns(column_names=['id', 'email'])))
+
+    catalog.alter_table('default.users', [SchemaChange.rename_column('email', 'contact')])
+    catalog.alter_table('default.users', [SchemaChange.drop_column('id')])
+    grants = catalog.permission_management().list_permissions(
+        ListPermissionsRequest(ResourceType.COLUMN, 'default', 'users')).elements
+    assert [grant.get_columns().get_column_names() for grant in grants] == [('contact',)]
+    with pytest.raises(BadRequestException):
+        catalog.alter_table('default.users', [SchemaChange.drop_column('contact')])
+
+
+def test_masks_of_two_query_principals_on_one_column_deny_the_query(rest):
+    catalog, server = rest
+    for principal, transform in (('analyst', NULL_MASK), ('auditor', STARS_MASK)):
+        catalog.policy_management().create_policy(
+            DataPolicy.column_mask(_table('users'), ColumnMask('email', transform), principal))
+
+    server.set_query_principals('analyst')
+    assert catalog.auth_table_query(USERS).column_masking == {'email': NULL_MASK}
+    server.set_query_principals('analyst', 'auditor')
+    with pytest.raises(TableNoPermissionException):
+        catalog.auth_table_query(USERS)