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 1e64b63083 [python] Add data policy endpoints to the REST mock server
(#10618)
1e64b63083 is described below
commit 1e64b63083950dd08cd6e93b85e2b7c37f8c8231
Author: Jiajia Li <[email protected]>
AuthorDate: Sun Oct 18 18:53:17 2026 +0800
[python] Add data policy endpoints to the REST mock server (#10618)
---
.../tests/rest/rest_permission_server_test.py | 373 ++++++++-
paimon-python/pypaimon/tests/rest/rest_server.py | 838 ++++++++++++++++++++-
.../pypaimon/tests/sql_permission_test.py | 110 ++-
3 files changed, 1283 insertions(+), 38 deletions(-)
diff --git a/paimon-python/pypaimon/tests/rest/rest_permission_server_test.py
b/paimon-python/pypaimon/tests/rest/rest_permission_server_test.py
index 51db1276d6..4d24b3a93f 100644
--- a/paimon-python/pypaimon/tests/rest/rest_permission_server_test.py
+++ b/paimon-python/pypaimon/tests/rest/rest_permission_server_test.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
+import json
+
import pyarrow as pa
import pytest
@@ -22,14 +24,21 @@ 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 IllegalArgumentError,
TableNoPermissionException
+ ForbiddenException,
NoSuchResourceException)
+from pypaimon.catalog.catalog_exception import (IllegalArgumentError,
TableNoPermissionException,
+ TableNotExistException)
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.policy_type import PolicyType
from pypaimon.management.resource_type import ResourceType
+from pypaimon.management.row_filter import RowFilter
+from pypaimon.schema.data_types import AtomicType
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
@@ -37,7 +46,8 @@ from pypaimon.tests.rest.rest_server import RESTCatalogServer
USERS = Identifier.create('default', 'users')
MEMBERS = Identifier.create('default', 'members')
COLUMNS = PermissionResource(ResourceType.COLUMN, 'default', 'users')
-ARROW_SCHEMA = pa.schema([('id', pa.int64()), ('email', pa.string()),
('phone', pa.string())])
+ARROW_SCHEMA = pa.schema([('id', pa.int64()), ('email', pa.string()),
('joined', pa.date32()), ('score', pa.float64()),
+ ('tags', pa.list_(pa.int32())), ('phone',
pa.string())])
def _schema(query_auth=True):
@@ -48,6 +58,7 @@ def _schema(query_auth=True):
def rest(tmp_path):
server = RESTCatalogServer(str(tmp_path), BearTokenAuthProvider('token'),
ConfigResponse(defaults={'prefix': 'mock'}),
'warehouse')
+ server.register_management_principal('analyst', 'auditor', 'reader',
'writer')
server.start()
try:
catalog = CatalogFactory.create({
@@ -186,7 +197,7 @@ def test_column_grants_follow_column_renames(rest):
table = catalog.get_table('default.users')
for schema in (table.table_schema, table.schema_manager.latest()):
- assert [field.name for field in schema.fields] == ['contact', 'phone']
+ assert [field.name for field in schema.fields] == ['contact',
'joined', 'score', 'tags', 'phone']
assert schema.options['query-auth.enabled'] == 'true'
assert _column_grants(catalog) == [(('contact',), None)]
@@ -217,3 +228,357 @@ def
test_tables_with_branches_cannot_have_column_grants(rest):
catalog.permission_management().grant_permission(grant)
with pytest.raises(IllegalArgumentError):
catalog.create_branch(USERS, 'b1')
+
+
+STARS_MASK = json.dumps({'name': 'CONCAT', 'inputs': ['***']})
+LOWER_MASK = json.dumps({'name': 'LOWER', 'inputs': [{'name': 'email', 'type':
'STRING'}]})
+
+
+def _row_filter(transform, column, column_type, function='EQUAL', literal=1):
+ return RowFilter(json.dumps({'kind': 'LEAF', 'function': function,
'literals': [literal], 'transform': {
+ 'name': transform, 'fieldRef': {'name': column, 'type':
column_type}}}))
+
+
+def _mask(column, transform, principal='analyst', table='users'):
+ return DataPolicy.column_mask(_table(table), ColumnMask(column,
transform), principal)
+
+
+def _policies(catalog, name):
+ policies =
catalog.policy_management().list_policies(ListPoliciesRequest(_table(name))).elements
+ return [(policy.get_resource().get_table(),
policy.get_column_mask().get_on_column()) for policy in policies]
+
+
+def test_table_policies_follow_rename_and_drop(rest):
+ catalog, _ = rest
+ catalog.policy_management().create_policy(_mask('email', STARS_MASK))
+
+ catalog.rename_table('default.users', 'default.members')
+ assert _policies(catalog, 'members') == [('members', 'email')]
+ catalog.create_table('default.users', _schema(), False)
+ assert _policies(catalog, 'users') == []
+
+ catalog.drop_table('default.members')
+ catalog.create_table('default.members', _schema(), False)
+ assert _policies(catalog, 'members') == []
+
+
+def test_dropping_a_database_removes_only_its_tables_and_policies(rest):
+ catalog, server = rest
+ for database in ('sales', 'sales.eu'):
+ catalog.create_database(database, False)
+ tables = [PermissionResource(ResourceType.TABLE, 'sales', 'eu.users'),
+ PermissionResource(ResourceType.TABLE, 'sales.eu', 'orders')]
+ for table in tables:
+ catalog.create_table(Identifier.create(table.get_database(),
table.get_table()), _schema(), False)
+ catalog.policy_management().create_policy(
+ DataPolicy.column_mask(table, ColumnMask('email', STARS_MASK),
'analyst'))
+
+ assert [catalog.list_tables(database) for database in ('sales',
'sales.eu')] == [['eu.users'], ['orders']]
+ with pytest.raises(ValueError):
+ catalog.drop_database('sales.eu', False, False)
+ catalog.drop_database('sales', False, True)
+ catalog.create_database('sales', False)
+ with pytest.raises(TableNotExistException):
+ catalog.get_table(Identifier.create('sales', 'eu.users'))
+ assert
len(catalog.policy_management().list_policies(ListPoliciesRequest(tables[1])).elements)
== 1
+ catalog.drop_database('sales.eu', False, True)
+ assert server.policy_store == {}
+
+
[email protected]('drop_database', [False, True])
+def test_a_recreated_table_starts_from_its_own_schema(rest, drop_database):
+ catalog, _ = rest
+ catalog.alter_table('default.users', [SchemaChange.set_option('a', 'b')])
+ if drop_database:
+ catalog.drop_database('default', False, True)
+ catalog.create_database('default', False)
+ else:
+ catalog.drop_table('default.users')
+ catalog.create_table('default.users',
Schema.from_pyarrow_schema(pa.schema([('id', pa.int64())])), False)
+ catalog.alter_table('default.users', [SchemaChange.set_option('a', 'b')])
+ assert [field.name for field in
catalog.get_table('default.users').table_schema.fields] == ['id']
+
+
[email protected]('branch', ['b1', 'missing'])
+def test_a_branch_is_not_dropped_as_a_table(rest, branch):
+ catalog, _ = rest
+ catalog.create_branch(USERS, 'b1')
+ with pytest.raises(TableNotExistException):
+ catalog.drop_table(Identifier('default', 'users', branch=branch))
+ catalog.alter_table('default.users', [SchemaChange.set_option('a', 'b')])
+ assert catalog.get_table('default.users').table_schema.options['a'] == 'b'
+ assert catalog.list_branches(USERS) == ['b1']
+
+
+def test_policies_are_managed_without_read_access(rest):
+ catalog, server = rest
+ server.add_no_permission_database('default')
+ server.add_no_permission_table(USERS)
+ catalog.policy_management().create_policy(_mask('email', STARS_MASK))
+ assert _policies(catalog, 'users') == [('users', 'email')]
+ catalog.policy_management().drop_policy(_table('users'),
PolicyType.COLUMN_MASKING, 'analyst', 'email', False)
+ assert _policies(catalog, 'users') == []
+
+
+def test_management_requests_check_principals_and_management_permission(rest):
+ catalog, server = rest
+ permissions, policies = catalog.permission_management(),
catalog.policy_management()
+ for request in (lambda:
permissions.grant_permission(PermissionAssignment(_table('users'), 'SELECT',
'nobody')),
+ lambda: permissions.revoke_permission(_table('users'),
'SELECT', 'nobody'),
+ lambda: permissions.list_permissions(
+ ListPermissionsRequest(ResourceType.TABLE, 'default',
'users', principal='nobody')),
+ lambda: policies.create_policy(_mask('email', STARS_MASK,
'nobody')),
+ lambda:
policies.list_policies(ListPoliciesRequest(_table('users'),
principal='nobody'))):
+ with pytest.raises(NoSuchResourceException) as error:
+ request()
+ assert error.value.resource_type == 'PRINCIPAL'
+ permissions.grant_permission(PermissionAssignment(_table('users'),
'SELECT', 'analyst'))
+ policies.create_policy(_mask('email', STARS_MASK))
+
+ server.deny_management_permission(_table('users'))
+ drop = (_table('users'), PolicyType.COLUMN_MASKING, 'analyst', 'email',
False)
+ for request in (lambda:
permissions.grant_permission(PermissionAssignment(_table('users'), 'SELECT',
'auditor')),
+ lambda: permissions.revoke_permission(_table('users'),
'SELECT', 'analyst'),
+ lambda: _grants(catalog, 'users'), lambda:
_policies(catalog, 'users'),
+ lambda: policies.create_policy(_mask('phone', STARS_MASK)),
+ lambda: policies.drop_policy(*drop)):
+ with pytest.raises(ForbiddenException):
+ request()
+ assert [len(server.permission_store), [len(table) for table in
server.policy_store.values()]] == [1, [1]]
+
+
+def test_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.policy_management().create_policy(_mask('email', STARS_MASK,
table='plain'))
+ for policy in (_mask('mobile', STARS_MASK),
+ DataPolicy.row_filter(_table('users'),
_row_filter('FIELD_REF', 'mobile', 'BIGINT'), 'analyst')):
+ with pytest.raises(BadRequestException):
+ catalog.policy_management().create_policy(policy)
+
+
[email protected]('policy, column, other', [
+ (_mask('email', STARS_MASK), 'email', 'id'),
+ (DataPolicy.row_filter(_table('users'), _row_filter('FIELD_REF', 'id',
'BIGINT'), 'analyst'), 'id', 'email'),
+ (DataPolicy.row_filter(_table('users'), _row_filter('YEAR', 'joined',
'DATE'), 'analyst'), 'joined', 'id'),
+])
+def test_schema_changes_cannot_orphan_policies(rest, policy, column, other):
+ catalog, _ = rest
+ catalog.policy_management().create_policy(policy)
+ schema = catalog.get_table('default.users').table_schema
+ column_type = next(field.type for field in schema.fields if field.name ==
column)
+ for changes in ([SchemaChange.set_option('query-auth.enabled', 'false')],
+ [SchemaChange.remove_option('query-auth.enabled')],
+ [SchemaChange.rename_column(column, 'renamed')],
+ [SchemaChange.drop_column(column)],
+ [SchemaChange.rename_column(column, 'renamed'),
SchemaChange.add_column(column, column_type)]):
+ with pytest.raises(BadRequestException):
+ catalog.alter_table('default.users', changes)
+
+ current = catalog.get_table('default.users').table_schema
+ assert (current.id, current.options) == (schema.id, schema.options)
+ policies =
catalog.policy_management().list_policies(ListPoliciesRequest(_table('users'))).elements
+ assert [stored.type() for stored in policies] == [policy.type()]
+ catalog.alter_table('default.users', [SchemaChange.rename_column(other,
'renamed')])
+
+
+def test_policy_expressions_must_fit_the_table_schema(rest):
+ catalog, _ = rest
+ policies = catalog.policy_management()
+ email, joined = {'name': 'email', 'type': 'STRING'}, {'name': 'joined',
'type': 'DATE'}
+ masks = [('email', {'name': 'NULL'}), ('email', {'name': 'TRIM', 'inputs':
[email]}),
+ ('email', {'name': 'SUBSTRING', 'inputs': [123, 1]}), ('email',
{'name': 'YEAR', 'fieldRef': joined}),
+ ('email', {'name': 'LOWER', 'inputs': [{'name': 'id', 'type':
'BIGINT'}]}),
+ ('joined', {'name': 'CAST', 'fieldRef': joined, 'type': 'DATE'}),
+ ('email', {'name': 'CAST', 'fieldRef': {'name': 'id', 'type':
'BIGINT'}, 'type': 'VARCHAR(2147483647)'}),
+ ('email', {'name': 'CONCAT', 'inputs': email}), ('email',
{'name': 'CONCAT', 'inputs': None}),
+ ('email', {'name': 'CONCAT', 'inputs': [{'name': 'email', 'type':
'INT'}, email]})]
+ row_filters = [RowFilter(json.dumps(predicate)) for predicate in (
+ None, {}, {'kind': 'COMPOUND', 'function': 'AND', 'children': []})] + [
+ _row_filter('FIELD_REF', 'id', 'BIGINT', literal=float('nan')),
+ _row_filter('YEAR', 'email', 'STRING'), _row_filter('FIELD_REF', 'id',
'BIGINT', literal='one'),
+ _row_filter('FIELD_REF', 'id', 'BIGINT', function='NOT_REAL'),
+ _row_filter('FIELD_REF', 'id', 'INTEGER'),
+ _row_filter('FIELD_REF', 'email', 'STRING', literal=1.5),
+ _row_filter('FIELD_REF', 'score', 'FLOAT', literal=1e11),
+ _row_filter('FIELD_REF', 'joined', 'DATE', literal='2023-02-29'),
+ RowFilter(json.dumps({'kind': 'LEAF', 'function': 'IS_NULL',
'literals': [], 'transform': {
+ 'name': 'CAST', 'fieldRef': {'name': 'score', 'type': 'DOUBLE'},
'type': 'TIMESTAMP(6)'}})),
+ RowFilter(json.dumps({'kind': 'LEAF', 'function': 'IS_NULL',
'literals': [], 'transform': {
+ 'name': 'CAST', 'fieldRef': email, 'type': {'type': 'ARRAY',
'element': {
+ 'type': 'VECTOR', 'element': 'FLOAT', 'length': 3}}}}))]
+ for policy in ([_mask(column, json.dumps(transform)) for column, transform
in masks]
+ + [DataPolicy.row_filter(_table('users'), row_filter,
'analyst') for row_filter in row_filters]):
+ with pytest.raises(BadRequestException):
+ policies.create_policy(policy)
+
+ policies.create_policy(_mask('email', LOWER_MASK))
+ policies.create_policy(_mask('email', json.dumps(
+ {'name': 'CAST', 'fieldRef': {'name': 'id', 'type': 'BIGINT'}, 'type':
'STRING'}), 'auditor'))
+ policies.create_policy(DataPolicy.row_filter(_table('users'),
_row_filter('YEAR', 'joined', 'DATE'), 'analyst'))
+ policies.create_policy(DataPolicy.row_filter(
+ _table('users'), _row_filter('FIELD_REF', 'joined', 'DATE',
literal='2024-02-29'), 'auditor'))
+ policies.create_policy(DataPolicy.row_filter(_table('users'),
RowFilter(json.dumps(
+ {'kind': 'LEAF', 'function': 'EQUAL', 'literals':
['2024-01-01T00:00'], 'transform': {
+ 'name': 'CAST', 'fieldRef': {'name': 'id', 'type': 'BIGINT'},
'type': 'TIMESTAMP(6)'}})), 'reader'))
+ policies.create_policy(DataPolicy.row_filter(
+ _table('users'), _row_filter('FIELD_REF', 'score', 'FLOAT',
literal=1.5), 'writer'))
+ for column, new_type in (('email', 'BIGINT'), ('joined', 'STRING')):
+ with pytest.raises(BadRequestException):
+ catalog.alter_table('default.users',
[SchemaChange.update_column_type(column, AtomicType(new_type))])
+
+
+def test_float_literals_block_widening_unless_exact(rest):
+ catalog, _ = rest
+ catalog.create_table('default.metrics', Schema.from_pyarrow_schema(
+ pa.schema([('v', pa.float32())]), options={'query-auth.enabled':
'true'}), False)
+ policies = catalog.policy_management()
+ widen = [SchemaChange.update_column_type('v', AtomicType('DOUBLE'))]
+
+ policies.create_policy(DataPolicy.row_filter(
+ _table('metrics'), _row_filter('FIELD_REF', 'v', 'FLOAT',
literal=1e11), 'analyst'))
+ with pytest.raises(BadRequestException):
+ catalog.alter_table('default.metrics', widen)
+
+ policies.drop_policy(_table('metrics'), PolicyType.ROW_FILTER, 'analyst',
None, False)
+ policies.create_policy(DataPolicy.row_filter(
+ _table('metrics'), _row_filter('FIELD_REF', 'v', 'FLOAT',
literal=1.5), 'analyst'))
+ catalog.alter_table('default.metrics', widen)
+ assert [field.type.type for field in
catalog.get_table('default.metrics').table_schema.fields] == ['DOUBLE']
+
+
+def test_instants_take_a_leap_second_only_at_the_end_of_a_day(rest):
+ catalog, _ = rest
+ catalog.create_table('default.events', Schema.from_pyarrow_schema(
+ pa.schema([('at', pa.timestamp('us', tz='UTC'))]),
options={'query-auth.enabled': 'true'}), False)
+ policies = catalog.policy_management()
+
+ with pytest.raises(BadRequestException):
+ policies.create_policy(DataPolicy.row_filter(_table('events'),
_row_filter(
+ 'FIELD_REF', 'at', 'TIMESTAMP_LTZ(6)',
literal='2024-01-01T12:34:60Z'), 'analyst'))
+ policies.create_policy(DataPolicy.row_filter(_table('events'), _row_filter(
+ 'FIELD_REF', 'at', 'TIMESTAMP_LTZ(6)',
literal='2016-12-31T23:59:60Z'), 'analyst'))
+ stored =
policies.list_policies(ListPoliciesRequest(_table('events'))).elements[0]
+ assert json.loads(stored.get_row_filter().get_predicate())['literals'] ==
['2016-12-31T23:59:59Z']
+
+
+def test_temporal_literals_keep_to_the_years_java_writes(rest):
+ catalog, _ = rest
+ catalog.create_table('default.events',
Schema.from_pyarrow_schema(pa.schema(
+ [('at', pa.timestamp('us', tz='UTC')), ('ts', pa.timestamp('us')),
('day', pa.date32())]),
+ options={'query-auth.enabled': 'true'}), False)
+ policies = catalog.policy_management()
+
+ for column, column_type, literal in (
+ ('at', 'TIMESTAMP_LTZ(6)', '+2024-01-01T00:00:00Z'), ('at',
'TIMESTAMP_LTZ(6)', '10000-01-01T00:00:00Z'),
+ ('at', 'TIMESTAMP_LTZ(6)', '-0000-01-01T00:00:00Z'),
+ ('at', 'TIMESTAMP_LTZ(6)', '+1000000000-01-01T00:00:00Z'),
+ ('ts', 'TIMESTAMP(6)', '-0000-01-01T00:00'), ('ts',
'TIMESTAMP(6)', '+5881580-07-12T00:00'),
+ ('day', 'DATE', '-0000-01-01')):
+ with pytest.raises(BadRequestException):
+ policies.create_policy(DataPolicy.row_filter(
+ _table('events'), _row_filter('FIELD_REF', column,
column_type, literal=literal), 'analyst'))
+ for principal, column, column_type, literal in (
+ ('analyst', 'at', 'TIMESTAMP_LTZ(6)', '+10000-01-01T00:00:00Z'),
+ ('auditor', 'ts', 'TIMESTAMP(6)', '+5881580-07-11T00:00')):
+ policies.create_policy(DataPolicy.row_filter(
+ _table('events'), _row_filter('FIELD_REF', column, column_type,
literal=literal), principal))
+ assert
sorted(json.loads(policy.get_row_filter().get_predicate())['literals'][0]
+ for policy in
policies.list_policies(ListPoliciesRequest(_table('events'))).elements) == [
+ '+10000-01-01T00:00:00Z', '+5881580-07-11T00:00']
+
+
+def test_policy_listings_are_paged(rest):
+ catalog, server = rest
+ principals = ['user{:03}'.format(i) for i in range(102)]
+ server.register_management_principal(*principals)
+ for principal in principals:
+ catalog.policy_management().create_policy(_mask('email', STARS_MASK,
principal))
+
+ first =
catalog.policy_management().list_policies(ListPoliciesRequest(_table('users'),
max_results=1))
+ second = catalog.policy_management().list_policies(
+ ListPoliciesRequest(_table('users'), page_token=first.next_page_token,
max_results=1000))
+ assert [len(first.elements), len(second.elements), second.next_page_token]
== [1, 101, None]
+ assert [policy.get_principal() for policy in first.elements +
second.elements] == principals
+
+
+def test_policies_are_stored_as_java_serializes_them(rest):
+ catalog, server = rest
+ leaves = [{'kind': 'LEAF', 'function': 'IN', 'literals': [1.9, 2],
+ 'transform': {'name': 'FIELD_REF', 'fieldRef': {'name': 'id',
'type': 'BIGINT'}}},
+ {'kind': 'LEAF', 'function': 'GREATER_THAN', 'literals':
['2024-02-29'],
+ 'transform': {'name': 'FIELD_REF', 'fieldRef': {'name':
'joined', 'type': 'DATE'}}},
+ {'kind': 'LEAF', 'function': 'EQUAL', 'literals': [1e10],
'transform': {
+ 'name': 'CAST', 'fieldRef': {'name': 'id', 'type':
'BIGINT'}, 'type': 'INT'}},
+ {'kind': 'LEAF', 'function': 'EQUAL', 'literals': ['2024-02-29'],
+ 'transform': {'name': 'FIELD_REF', 'fieldRef': {'name': 'id',
'type': 'DATE'}}}]
+ leaves[1] = {'kind': 'COMPOUND', 'function': 'OR', 'children': [leaves[1]]}
+
catalog.policy_management().create_policy(DataPolicy.row_filter(_table('users'),
RowFilter(json.dumps(
+ {'kind': 'COMPOUND', 'function': 'AND', 'children': leaves})),
'analyst'))
+ catalog.policy_management().create_policy(_mask('email', LOWER_MASK))
+ catalog.policy_management().create_policy(_mask('tags', json.dumps(
+ {'name': 'CAST', 'fieldRef': {'name': 'email', 'type': 'STRING'},
+ 'type': {'type': 'ARRAY', 'element': 'INT'}})))
+
+ server.set_query_principals('analyst')
+ auth = catalog.auth_table_query(USERS)
+ assert auth.filter == [
+ '{"kind":"COMPOUND","function":"AND","children":['
+
'{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":0,"name":"id","type":"BIGINT"}},'
+ '"function":"IN","literals":[1,2]},'
+
'{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":2,"name":"joined","type":"DATE"}},'
+ '"function":"GREATER_THAN","literals":["2024-02-29"]},'
+
'{"kind":"LEAF","transform":{"name":"CAST","fieldRef":{"index":0,"name":"id","type":"BIGINT"},"type":"INT"},'
+ '"function":"EQUAL","literals":[2147483647]},'
+
'{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":0,"name":"id","type":"BIGINT"}},'
+ '"function":"EQUAL","literals":[19782]}]}']
+ assert auth.column_masking == {
+ 'email':
'{"name":"LOWER","inputs":[{"index":1,"name":"email","type":"STRING"}]}',
+ 'tags':
'{"name":"CAST","fieldRef":{"index":1,"name":"email","type":"STRING"},'
+ '"type":{"type":"ARRAY","element":"INT"}}'}
+
+
+def test_query_auth_returns_policies_for_the_current_schema(rest):
+ catalog, server = rest
+ catalog.policy_management().create_policy(
+ DataPolicy.row_filter(_table('users'), _row_filter('FIELD_REF', 'id',
'BIGINT'), 'analyst'))
+ catalog.policy_management().create_policy(
+ DataPolicy.row_filter(_table('users'), _row_filter('FIELD_REF',
'score', 'DOUBLE', literal=1.5), 'auditor'))
+ for column, column_type in (('id', 'DECIMAL(20, 0)'), ('score', 'STRING')):
+ with pytest.raises(BadRequestException):
+ catalog.alter_table('default.users',
[SchemaChange.update_column_type(column, AtomicType(column_type))])
+ catalog.alter_table('default.users',
[SchemaChange.update_column_type('id', AtomicType('STRING'))])
+
+ server.set_query_principals('analyst')
+ assert catalog.auth_table_query(USERS).filter == [
+
'{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":0,"name":"id","type":"STRING"}},'
+ '"function":"EQUAL","literals":["1"]}']
+
+
+def test_masks_of_two_query_principals_on_one_column_deny_the_query(rest):
+ catalog, server = rest
+ for principal, transform in (('analyst', STARS_MASK), ('auditor',
LOWER_MASK)):
+ catalog.policy_management().create_policy(_mask('email', transform,
principal))
+
+ server.set_query_principals('analyst')
+ assert catalog.auth_table_query(USERS).column_masking == {'email':
'{"name":"CONCAT","inputs":["***"]}'}
+ server.set_query_principals('analyst', 'auditor')
+ with pytest.raises(TableNoPermissionException):
+ catalog.auth_table_query(USERS)
+
+
+def test_tables_with_branches_cannot_have_policies(rest):
+ catalog, _ = rest
+ with pytest.raises(NoSuchResourceException):
+ catalog.policy_management().create_policy(_mask('email', STARS_MASK,
table='users$branch_missing'))
+ with pytest.raises(TableNotExistException):
+ catalog.auth_table_query(Identifier('default', 'users',
branch='missing'))
+ catalog.create_branch(USERS, 'b1')
+ with pytest.raises(BadRequestException):
+ catalog.policy_management().create_policy(_mask('email', STARS_MASK))
+
+ catalog.drop_branch(USERS, 'b1')
+ catalog.policy_management().create_policy(_mask('email', STARS_MASK))
+ with pytest.raises(IllegalArgumentError):
+ catalog.create_branch(USERS, 'b1')
diff --git a/paimon-python/pypaimon/tests/rest/rest_server.py
b/paimon-python/pypaimon/tests/rest/rest_server.py
index e61ea4388b..33bd32404f 100755
--- a/paimon-python/pypaimon/tests/rest/rest_server.py
+++ b/paimon-python/pypaimon/tests/rest/rest_server.py
@@ -15,11 +15,16 @@
# specific language governing permissions and limitations
# under the License.
+import base64
+import decimal
import logging
+import math
import re
+import struct
import threading
import time
import uuid
+from collections import namedtuple
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
import json
@@ -34,7 +39,8 @@ from pypaimon.api.api_request import (AlterDatabaseRequest,
AlterTableRequest,
CreateDatabaseRequest,
CreatePartitionsRequest,
CreateTableRequest, CreateTagRequest,
- GrantPermissionRequest,
+ DropPolicyRequest,
+ GrantPermissionRequest, PolicyRequest,
RenameTableRequest,
RevokePermissionRequest)
from pypaimon.api.api_response import (ConfigResponse,
CreatePartitionsResponse,
@@ -47,7 +53,8 @@ from pypaimon.api.api_response import (ConfigResponse,
CreatePartitionsResponse,
ListFunctionsGloballyResponse,
ListFunctionsResponse,
ListPartitionsResponse,
- ListPermissionsResponse,
ListTablesResponse,
+ ListPermissionsResponse,
+ ListPoliciesResponse,
ListTablesResponse,
ListTagsResponse,
PagedList, Partition,
RESTResponse, ErrorResponse)
@@ -66,18 +73,26 @@ from pypaimon.catalog.catalog_exception import
(BranchAlreadyExistException,
DefinitionNotExistException,
TagNotExistException,
TagAlreadyExistException)
+from pypaimon.casting import data_type_casts as casts
from pypaimon.catalog.rest.table_metadata import TableMetadata
from pypaimon.common.identifier import Identifier
from pypaimon.api.typedef import RESTAuthParameter
from pypaimon.common.json_util import JSON
from pypaimon.common.options.core_options import CoreOptions
+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.permission_access import PermissionAccess
from pypaimon.management.permission_assignment import (PermissionAssignment,
_INSTANT,
- _days_from_civil)
+ _days_from_civil,
_days_in_month,
+ _parse_instant as
_instant_nanos)
from pypaimon.management.permission_columns import PermissionColumns
from pypaimon.management.permission_resource import PermissionResource
from pypaimon.management.resource_type import ResourceType
+from pypaimon.management.row_filter import RowFilter
from pypaimon import Schema
+from pypaimon.schema.data_types import (ArrayType, AtomicType, DataField,
DataType, DataTypeParser, MapType,
+ MultisetType, RowType, VectorType)
from pypaimon.schema.schema_change import Actions, SchemaChange
from pypaimon.schema.schema_manager import SchemaManager
from pypaimon.schema.table_schema import TableSchema
@@ -227,6 +242,599 @@ def _resource_name(resource: PermissionResource) -> str:
return ".".join(part for part in (resource.get_database(), name) if part)
or "catalog"
+def _policy_identity(policy: DataPolicy) -> Tuple[str, str, Optional[str]]:
+ mask = policy.get_column_mask()
+ return policy.type().name, policy.get_principal(), None if mask is None
else mask.get_on_column()
+
+
+def _policy_key(policy: DataPolicy) -> str:
+ return _cursor_key(*_policy_identity(policy))
+
+
+def _policy_name(policy_type: str, principal: str, column: Optional[str]) ->
str:
+ return ":".join(part for part in (policy_type, principal, column) if part
is not None)
+
+
+def _field_refs(transform: Dict[str, Any]) -> List[Dict[str, Any]]:
+ refs = list(transform.get("inputs") or []) + [transform.get("fieldRef")]
+ return [ref for ref in refs if isinstance(ref, dict)]
+
+
+def _leaves(predicate: Dict[str, Any]) -> List[Dict[str, Any]]:
+ if predicate["kind"] == "COMPOUND":
+ children = predicate.get("children")
+ if predicate.get("function") not in ("AND", "OR") or not children or
None in children:
+ raise ValueError("Compound predicate needs AND or OR and at least
one predicate.")
+ return [leaf for child in children for leaf in _leaves(child)]
+ if predicate["kind"] != "LEAF" or not
isinstance(predicate.get("transform"), dict):
+ raise ValueError("Unknown predicate {}.".format(json.dumps(predicate)))
+ return [predicate]
+
+
+def _reject_constant(token: str) -> None:
+ raise ValueError("Non-standard JSON token {}.".format(token))
+
+
+def _policy_json(text: str) -> Any:
+ return json.loads(text, parse_constant=_reject_constant)
+
+
+def _policy_transforms(policy: DataPolicy) -> List[Tuple[Dict[str, Any],
Optional[Dict[str, Any]]]]:
+ mask = policy.get_column_mask()
+ if mask is not None:
+ return [(_policy_json(mask.get_transform()), None)]
+ return [(leaf["transform"], leaf) for leaf in
_leaves(_policy_json(policy.get_row_filter().get_predicate()))]
+
+
+def _policy_columns(policy: DataPolicy) -> set:
+ mask = policy.get_column_mask()
+ columns = {ref["name"] for transform, _ in _policy_transforms(policy) for
ref in _field_refs(transform)}
+ return columns | ({mask.get_on_column()} if mask is not None else set())
+
+
+_DATE_PARTS = ("YEAR", "QUARTER", "MONTH", "WEEK", "YEAR_OF_WEEK", "DAY",
"DAY_OF_WEEK", "ISO_DAY_OF_WEEK",
+ "WEEKDAY", "DAY_OF_YEAR", "HOUR", "MINUTE", "SECOND")
+_TRANSFORM_INPUTS = {
+ "CONCAT": ("s", 0, None, "STRING"), "CONCAT_WS": ("s", 2, None, "STRING"),
+ "UPPER": ("s", 1, 1, "STRING"), "LOWER": ("s", 1, 1, "STRING"), "TRIM":
("s", 1, 2, "STRING"),
+ "TRANSLATE": ("s", 3, 3, "STRING"), "SUBSTRING": ("ra", 2, 3, "STRING"),
+ "OVERLAY": ("ssii", 3, 4, "STRING"), "PAD": ("sis", 3, 3, "STRING"),
+ "LENGTH": ("s", 1, 1, "INT"), "BIT_LENGTH": ("s", 1, 1, "INT"),
+ "DATE_ADD": ("di", 2, 2, "DATE"), "DATE_DIFF": ("dd", 2, 2, "INT"),
"DATE_TRUNC": ("ds", 2, 2, "DATE"),
+}
+_TRANSFORM_FLAGS = {"TRIM": ("trimFlag", ("LEADING", "TRAILING", "BOTH")),
"PAD": ("direction", ("LEFT", "RIGHT"))}
+_INPUT_ROOTS = {"s": casts.CHARACTER_STRING, "i": casts.INTEGER_NUMERIC, "d":
{casts.DATE}}
+_TYPE_DEFAULTS = {casts.CHAR: (1,), casts.VARCHAR: (1,), casts.BINARY: (1,),
casts.VARBINARY: (1,),
+ casts.DECIMAL: (10, 0), casts.TIME: (0,), casts.TIMESTAMP:
(6,), casts.TIMESTAMP_LTZ: (6,)}
+_LEAF_FUNCTIONS = {
+ "EQUAL", "NOT_EQUAL", "LESS_THAN", "LESS_OR_EQUAL", "GREATER_THAN",
"GREATER_OR_EQUAL", "IS_NULL",
+ "IS_NOT_NULL", "STARTS_WITH", "NOT_STARTS_WITH", "ENDS_WITH",
"NOT_ENDS_WITH", "CONTAINS", "NOT_CONTAINS",
+ "ARRAY_CONTAINS", "ARRAYS_OVERLAP", "ARRAY_CONTAINS_ALL", "LIKE",
"NOT_LIKE", "IN", "NOT_IN", "BETWEEN",
+ "NOT_BETWEEN", "TRUE", "FALSE"}
+_ARRAY_FUNCTIONS = ("ARRAY_CONTAINS", "ARRAYS_OVERLAP", "ARRAY_CONTAINS_ALL")
+_PRIMITIVES = casts.INTEGER_NUMERIC | {casts.FLOAT, casts.DOUBLE}
+_CAST_RULES = (
+ (casts.CHARACTER_STRING, casts.CHARACTER_STRING | casts.BINARY_STRING |
_PRIMITIVES | casts.TIMESTAMP_FAMILY
+ | {casts.BOOLEAN, casts.DATE, casts.DECIMAL, casts.TIME, casts.ARRAY,
casts.MAP, casts.ROW}),
+ (casts.BINARY_STRING, casts.BINARY_STRING | casts.CHARACTER_STRING),
+ ({casts.BOOLEAN}, casts.NUMERIC | casts.CHARACTER_STRING),
+ ({casts.DATE}, casts.CHARACTER_STRING | casts.TIMESTAMP_FAMILY),
+ ({casts.DECIMAL}, _PRIMITIVES | {casts.DECIMAL}),
+ (_PRIMITIVES, _PRIMITIVES | {casts.DECIMAL}),
+ (casts.INTEGER_NUMERIC, {casts.BOOLEAN}),
+ (casts.NUMERIC, casts.CHARACTER_STRING),
+ ({casts.INTEGER, casts.BIGINT}, casts.TIMESTAMP_FAMILY),
+ ({casts.ARRAY, casts.MAP, casts.ROW}, casts.CHARACTER_STRING),
+ (casts.TIMESTAMP_FAMILY, casts.CHARACTER_STRING | casts.TIMESTAMP_FAMILY
+ | {casts.DATE, casts.TIME, casts.BIGINT, casts.INTEGER}),
+ ({casts.TIME}, casts.CHARACTER_STRING | casts.TIMESTAMP_FAMILY),
+)
+_DECIMAL_LITERAL =
re.compile(r"[+-]?(?:[0-9]+\.?[0-9]*|\.[0-9]+)(?:[eE][+-]?[0-9]+)?")
+_LOCAL_DATE =
re.compile(r"([0-9]{4}|\+[0-9]{5,9}|-[0-9]{4,9})-([0-9]{2})-([0-9]{2})")
+_LOCAL_TIME =
re.compile(r"([0-9]{2}):([0-9]{2})(?::([0-9]{2})(?:\.([0-9]{0,9}))?)?")
+_INTEGER_BITS = {casts.TINYINT: 8, casts.SMALLINT: 16, casts.INTEGER: 32,
casts.BIGINT: 64}
+
+
+def _type_root(data_type: Optional[DataType]) -> Optional[str]:
+ if isinstance(data_type, AtomicType) and "LOCAL TIME ZONE" in
data_type.type.upper():
+ return casts.TIMESTAMP_LTZ
+ return casts._root(data_type)
+
+
+def _logical_type(data_type: Optional[DataType]) -> Any:
+ if isinstance(data_type, RowType):
+ return "ROW", tuple((field.id, field.name, _logical_type(field.type),
field.description, field.default_value)
+ for field in data_type.fields), data_type.nullable
+ if isinstance(data_type, MapType):
+ return "MAP", _logical_type(data_type.key),
_logical_type(data_type.value), data_type.nullable
+ if isinstance(data_type, (ArrayType, MultisetType, VectorType)):
+ return type(data_type).__name__, _logical_type(data_type.element),
getattr(data_type, "length", None), \
+ data_type.nullable
+ if not isinstance(data_type, AtomicType):
+ return data_type
+ spelling = data_type.type.upper()
+ root = _type_root(data_type)
+ found = (2147483647,) if spelling in ("STRING", "BYTES") else
tuple(map(int, re.findall(r"\d+", spelling)))
+ return root, found + _TYPE_DEFAULTS.get(root, ())[len(found):],
data_type.nullable
+
+
+def _castable(source: DataType, target: DataType) -> bool:
+ if not any(_type_root(source) in inputs and _type_root(target) in targets
for inputs, targets in _CAST_RULES):
+ return False
+ if _type_root(source) not in casts.CHARACTER_STRING or not
isinstance(target, (ArrayType, MapType, RowType)):
+ return True
+ children = [field.type for field in target.fields] if isinstance(target,
RowType) else \
+ [target.key, target.value] if isinstance(target, MapType) else
[target.element]
+ return all(_castable(source, child) for child in children)
+
+
+def _java_integer(value: Any, bits: int) -> int:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError("Not a number: {!r}".format(value))
+ if isinstance(value, float):
+ limit = 2 ** (63 if bits == 64 else 31)
+ value = 0 if math.isnan(value) else int(max(-limit, min(limit - 1,
value)))
+ value &= (1 << bits) - 1
+ return value - (1 << bits) if value >> (bits - 1) else value
+
+
+class _JavaFloat(float):
+ pass
+
+
+def _float32(value: float) -> float:
+ return struct.unpack("f", struct.pack("f", value))[0]
+
+
+def _java_float(value: Any, single: bool) -> float:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError("Not a number: {!r}".format(value))
+ try:
+ value = float(value)
+ return _JavaFloat(_float32(value)) if single else value
+ except OverflowError:
+ return (_JavaFloat if single else float)(math.inf if value > 0 else
-math.inf)
+
+
+def _java_number_string(value: float, single: bool = False) -> str:
+ if math.isnan(value) or math.isinf(value):
+ return "NaN" if math.isnan(value) else "Infinity" if value > 0 else
"-Infinity"
+ if value == 0:
+ return "-0.0" if math.copysign(1, value) < 0 else "0.0"
+ fits = (lambda text: _float32(float(text)) == value) if single else
(lambda text: float(text) == value)
+ digits = next(digits for digits in range(1, 18) if
fits("{:.{}e}".format(value, digits - 1)))
+ number = decimal.Decimal("{:.{}e}".format(value, max(digits, 2) -
1)).normalize()
+ if 1e-3 <= abs(value) < 1e7:
+ text = "{:f}".format(number)
+ return text if "." in text else text + ".0"
+ sign, digits, exponent = number.as_tuple()
+ mantissa = "".join(map(str, digits))
+ return "{}{}.{}E{}".format("-" if sign else "", mantissa[0], mantissa[1:]
or "0", exponent + len(digits) - 1)
+
+
+class _UnsupportedConversion(ValueError):
+ pass
+
+
+def _java_string(value: Any) -> str:
+ if isinstance(value, (float, list, dict)):
+ raise _UnsupportedConversion("Policy literal {} cannot be compared as
a string.".format(json.dumps(value)))
+ if isinstance(value, _Binary):
+ return value.text
+ if isinstance(value, _Decimal):
+ return "{:f}".format(value.value)
+ if isinstance(value, _Timestamp):
+ return _java_local_date_time(value)
+ if isinstance(value, str):
+ return value
+ if isinstance(value, bool) or value is None:
+ return {True: "true", False: "false", None: "null"}[value]
+ if isinstance(value, float):
+ return _java_number_string(value, isinstance(value, _JavaFloat))
+ if isinstance(value, list):
+ return "[" + ", ".join(map(_java_string, value)) + "]"
+ if isinstance(value, dict):
+ return "{" + ", ".join("{}={}".format(key, _java_string(item)) for
key, item in value.items()) + "}"
+ return str(value)
+
+
+_Binary = namedtuple("_Binary", "text")
+_Decimal = namedtuple("_Decimal", "value precision")
+_Timestamp = namedtuple("_Timestamp", "millisecond nano_of_millisecond")
+
+
+def _java_decimal(value: Any, literal_type: DataType) -> Optional[_Decimal]:
+ if not isinstance(value, str) or _DECIMAL_LITERAL.fullmatch(value) is None:
+ raise ValueError("Not a decimal: {!r}".format(value))
+ precision, scale = _logical_type(literal_type)[1]
+ with decimal.localcontext() as context:
+ context.prec, context.Emax, context.Emin = 10 ** 6, decimal.MAX_EMAX,
decimal.MIN_EMIN
+ try:
+ value =
decimal.Decimal(value).quantize(decimal.Decimal(1).scaleb(-scale),
rounding=decimal.ROUND_HALF_UP)
+ except decimal.InvalidOperation:
+ return None
+ if len(value.as_tuple().digits) > precision:
+ return None
+ return _Decimal(abs(value) if value.is_zero() else value, precision)
+
+
+def _java_date(year: int, month: int, day: int) -> str:
+ return "{}{:04d}-{:02d}-{:02d}".format("-" if year < 0 else "+" if year >
9999 else "", abs(year), month, day)
+
+
+def _java_time(hour: int, minute: int, second: int, nanos: int, seconds: bool
= False) -> str:
+ text = "{:02d}:{:02d}".format(hour, minute)
+ if second or nanos or seconds:
+ text += ":{:02d}".format(second)
+ if nanos:
+ fraction = "{:09d}".format(nanos)
+ text += "." + (fraction[:3] if nanos % 10 ** 6 == 0 else fraction[:6]
if nanos % 1000 == 0 else fraction)
+ return text
+
+
+def _parse_date(text: str) -> Tuple[int, int, int]:
+ match = _LOCAL_DATE.fullmatch(text)
+ year, month, day = map(int, match.groups()) if match else (0, 0, 0)
+ if not 1 <= month <= 12 or not 1 <= day <= _days_in_month(year, month) or
(year == 0 and text[0] == "-"):
+ raise ValueError("Not a date: {!r}".format(text))
+ return year, month, day
+
+
+def _parse_time(text: str) -> Tuple[int, int, int, int]:
+ match = _LOCAL_TIME.fullmatch(text)
+ if match is None or int(match.group(1)) > 23 or int(match.group(2)) > 59
or int(match.group(3) or 0) > 59:
+ raise ValueError("Not a time: {!r}".format(text))
+ return int(match.group(1)), int(match.group(2)), int(match.group(3) or 0),
int((match.group(4) or "").ljust(9, "0"))
+
+
+def _civil_from_days(days: int) -> Tuple[int, int, int]:
+ days += 719468
+ era = days // 146097
+ day_of_era = days - era * 146097
+ year_of_era = (day_of_era - day_of_era // 1460 + day_of_era // 36524 -
day_of_era // 146096) // 365
+ day_of_year = day_of_era - (365 * year_of_era + year_of_era // 4 -
year_of_era // 100)
+ shifted_month = (5 * day_of_year + 2) // 153
+ month = shifted_month + 3 if shifted_month < 10 else shifted_month - 9
+ return year_of_era + era * 400 + (month <= 2), month, day_of_year - (153 *
shifted_month + 2) // 5 + 1
+
+
+def _timestamp(days: int, hour: int, minute: int, second: int, nanos: int) ->
_Timestamp:
+ nano_of_day = ((hour * 60 + minute) * 60 + second) * 10 ** 9 + nanos
+ millisecond = days * 86400000 + nano_of_day // 10 ** 6
+ if not -2 ** 63 <= millisecond < 2 ** 63:
+ raise _UnsupportedConversion("Timestamp literal overflows the epoch
milliseconds Java keeps in a long.")
+ return _Timestamp(millisecond, nano_of_day % 10 ** 6)
+
+
+def _parse_instant(text: str) -> _Timestamp:
+ nanos = _instant_nanos(text)
+ if nanos is None:
+ raise ValueError("Not an instant: {!r}".format(text))
+ sign, digits, month, day, hour, minute, second, _ =
_INSTANT.fullmatch(text).groups()
+ hour, minute, second = int(hour), int(minute), int(second)
+ next_day = hour == 24
+ hour, second = 0 if next_day else hour, 59 if (hour, minute, second) ==
(23, 59, 60) else second
+ return _timestamp(_days_from_civil(int(sign + digits), int(month),
int(day)) + next_day, hour, minute, second,
+ nanos)
+
+
+def _date_time_parts(timestamp: _Timestamp) -> Tuple[str, int, int, int, int]:
+ days, millis = divmod(timestamp.millisecond, 86400000)
+ seconds, nanos = divmod(millis * 10 ** 6 + timestamp.nano_of_millisecond,
10 ** 9)
+ return _java_date(*_civil_from_days(days)), seconds // 3600, seconds // 60
% 60, seconds % 60, nanos
+
+
+def _java_local_date_time(timestamp: _Timestamp) -> str:
+ if not -2 ** 31 <= timestamp.millisecond // 86400000 < 2 ** 31:
+ raise _UnsupportedConversion("Timestamp literal overflows the epoch
days Java keeps in an int.")
+ date, hour, minute, second, nanos = _date_time_parts(timestamp)
+ return date + "T" + _java_time(hour, minute, second, nanos)
+
+
+def _java_instant(timestamp: _Timestamp) -> str:
+ date, hour, minute, second, nanos = _date_time_parts(timestamp)
+ return date + "T" + _java_time(hour, minute, second, nanos, seconds=True)
+ "Z"
+
+
+def _internal_literal(literal_type: DataType, literal: Any) -> Any:
+ root = _type_root(literal_type)
+ if literal is None or root == casts.BOOLEAN:
+ return literal
+ if root in _INTEGER_BITS:
+ return _java_integer(literal, _INTEGER_BITS[root])
+ if root in (casts.FLOAT, casts.DOUBLE):
+ return _java_float(literal, root == casts.FLOAT)
+ if root in casts.CHARACTER_STRING:
+ return _Binary(_java_string(literal))
+ if not isinstance(literal, str):
+ raise ValueError("{} literal must be a string: {!r}".format(root,
literal))
+ if root == casts.DECIMAL:
+ return _java_decimal(literal, literal_type)
+ if root == casts.DATE:
+ return _days_from_civil(*_parse_date(literal))
+ if root == casts.TIME:
+ hour, minute, second, nanos = _parse_time(literal)
+ return ((hour * 60 + minute) * 60 + second) * 1000 + nanos // 10 ** 6
+ if root == casts.TIMESTAMP:
+ parts = re.split("[Tt]", literal, maxsplit=1)
+ if len(parts) != 2:
+ raise ValueError("Not a timestamp: {!r}".format(literal))
+ return _timestamp(_days_from_civil(*_parse_date(parts[0])),
*_parse_time(parts[1]))
+ if root == casts.TIMESTAMP_LTZ:
+ return _parse_instant(literal)
+ raise ValueError("Unsupported literal type {}.".format(literal_type))
+
+
+def _json_literal(literal_type: DataType, value: Any) -> Any:
+ root = _type_root(literal_type)
+ single = isinstance(value, _JavaFloat)
+ if isinstance(value, float) and root != (casts.FLOAT if single else
casts.DOUBLE) \
+ and (root in _INTEGER_BITS or root in (casts.FLOAT, casts.DOUBLE))
\
+ and decimal.Decimal(_java_number_string(value, single)) !=
decimal.Decimal(value):
+ raise _UnsupportedConversion("Policy literal {} cannot be written as
{} because JDK versions write it "
+
"differently.".format(_java_number_string(value, single), literal_type))
+ if value is None or root == casts.BOOLEAN or root in _INTEGER_BITS or root
in (casts.FLOAT, casts.DOUBLE):
+ return value
+ if root in casts.CHARACTER_STRING:
+ return _java_string(value)
+ if root in (casts.DATE, casts.TIME):
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise RuntimeError("{!r} cannot be cast to a
Number.".format(value))
+ number = _java_integer(value, 32)
+ if root == casts.DATE:
+ return _java_date(*_civil_from_days(number))
+ if not 0 <= number < 86400000:
+ raise RuntimeError("Invalid time of day {}.".format(number))
+ return _java_time(number // 3600000, number // 60000 % 60, number //
1000 % 60, number % 1000 * 10 ** 6)
+ if root == casts.DECIMAL and isinstance(value, _Decimal):
+ return _java_string(value)
+ if root in (casts.TIMESTAMP, casts.TIMESTAMP_LTZ) and isinstance(value,
_Timestamp):
+ return _java_local_date_time(value) if root == casts.TIMESTAMP else
_java_instant(value)
+ raise RuntimeError("{!r} cannot be written as {}.".format(value,
literal_type))
+
+
+def _literal_type(function: str, output: DataType, error: type = ValueError)
-> DataType:
+ if function not in _ARRAY_FUNCTIONS:
+ return output
+ if not isinstance(output, ArrayType):
+ raise error("{} needs an ARRAY field.".format(function))
+ return output.element
+
+
+def _leaf_literals(leaf: Dict[str, Any], output: DataType) ->
Optional[List[Any]]:
+ function, literals = leaf.get("function"), leaf.get("literals")
+ if function not in _LEAF_FUNCTIONS or not isinstance(literals, (list,
type(None))):
+ raise ValueError("Unknown leaf function {} or literals
{}.".format(function, literals))
+ literal_type = _literal_type(function, output)
+ return None if literals is None else [_internal_literal(literal_type,
literal) for literal in literals]
+
+
+def _ref_type(ref: Dict[str, Any], columns: Optional[Dict[str, DataType]]) ->
DataType:
+ return _declared_type(ref["type"]) if columns is None else
columns[ref["name"]]
+
+
+def _input_fits(kind: str, value: Any, columns: Optional[Dict[str, DataType]])
-> bool:
+ if value is None:
+ return True
+ if isinstance(value, dict):
+ return kind in "ra" or _type_root(_ref_type(value, columns)) in
_INPUT_ROOTS[kind]
+ if isinstance(value, str):
+ return kind in "sra"
+ if isinstance(value, bool) or not isinstance(value, int):
+ return False
+ return kind in "ia" or (kind == "d" and -2 ** 31 <= value < 2 ** 31)
+
+
+def _transform_type(transform: Dict[str, Any],
+ columns: Optional[Dict[str, DataType]] = None) ->
Optional[DataType]:
+ name = transform["name"]
+ if "fieldRef" in transform:
+ source = _ref_type(transform["fieldRef"], columns)
+ if name == "CAST":
+ target = _declared_type(transform["type"])
+ return target if columns is not None or _castable(source, target)
else None
+ if name in _DATE_PARTS:
+ return AtomicType("INT") if _type_root(source) in (casts.DATE,
casts.TIMESTAMP) else None
+ for part in transform.get("path") or []:
+ source = next((field.type for field in getattr(source, "fields",
()) if field.name == part), None)
+ nested = name == "NESTED_FIELD_REF" and transform.get("path")
+ return source if name == "FIELD_REF" or nested else None
+ if name == "NULL":
+ return AtomicType("BOOLEAN")
+ if name in _TRANSFORM_FLAGS:
+ flag, values = _TRANSFORM_FLAGS[name]
+ if transform.get(flag) not in values:
+ return None
+ kinds, low, high, output = _TRANSFORM_INPUTS.get(name, ("", 1, 0, None))
+ inputs = transform.get("inputs")
+ if not isinstance(inputs, list) or not low <= len(inputs) <= (len(inputs)
if high is None else high):
+ return None
+ fits = all(_input_fits(kinds[min(position, len(kinds) - 1)], value,
columns)
+ for position, value in enumerate(inputs))
+ return AtomicType(output) if fits else None
+
+
+def _policy_type_error(policy: DataPolicy, schema: TableSchema) ->
Optional[str]:
+ columns = {field.name: field.type for field in schema.fields}
+ output = None
+ for transform, leaf in _policy_transforms(policy):
+ declared_types = _declared_types(transform)
+ declared = _transform_type(transform)
+ output = _transform_type(transform, columns)
+ if declared is None or output is None:
+ return "Policy transform {} does not fit the table
schema.".format(transform["name"])
+ if transform["name"] == "NESTED_FIELD_REF":
+ name, path = transform["fieldRef"]["name"], transform["path"]
+ if _path_ids(declared_types[name], path) !=
_path_ids(columns[name], path):
+ return "Nested field {}.{} changed identity.".format(name,
".".join(path))
+ try:
+ if leaf is not None:
+ _leaf_literals(leaf, declared)
+ except ValueError as e:
+ return "Row filter function {} does not accept {}: {}".format(
+ leaf.get("function"), leaf.get("literals"), e)
+ mask = policy.get_column_mask()
+ if mask is not None and _logical_type(output) !=
_logical_type(columns[mask.get_on_column()]):
+ return "Column mask output type {} does not match target column {}
type {}.".format(
+ output, mask.get_on_column(), columns[mask.get_on_column()])
+ return None
+
+
+def _java_spelling(value: Any) -> Any:
+ if isinstance(value, dict):
+ return {key: _java_spelling(item) for key, item in value.items() if
key != "nullable" and item is not None}
+ if isinstance(value, list):
+ return [_java_spelling(item) for item in value]
+ if isinstance(value, str):
+ return re.sub(r"^TIMESTAMP_LTZ\(([0-9]+)\)", r"TIMESTAMP(\1) WITH
LOCAL TIME ZONE", value)
+ return value
+
+
+_TYPE_RANGES = {casts.CHAR: ((1, 2147483647),), casts.VARCHAR: ((1,
2147483647),),
+ casts.BINARY: ((1, 2147483647),), casts.VARBINARY: ((1,
2147483647),), casts.DECIMAL: ((1, 38),),
+ casts.TIME: ((0, 9),), casts.TIMESTAMP: ((0, 9),),
casts.TIMESTAMP_LTZ: ((0, 9),)}
+
+
+def _valid_type(data_type: DataType) -> bool:
+ if not isinstance(data_type, AtomicType):
+ children = [field.type for field in getattr(data_type, "fields", ())]
+ [
+ getattr(data_type, name) for name in ("element", "key", "value")
if hasattr(data_type, name)]
+ return all(map(_valid_type, children))
+ root, parameters, _ = _logical_type(data_type)
+ ranges = _TYPE_RANGES.get(root, ())
+ if root == casts.DECIMAL and len(parameters) == 2 and not 0 <=
parameters[1] <= parameters[0]:
+ return False
+ return all(low <= value <= high for value, (low, high) in zip(parameters,
ranges))
+
+
+def _declared_type(value: Any) -> DataType:
+ data_type = DataTypeParser.parse_data_type(value)
+ if _java_spelling(value) != _java_type(data_type) or not
_valid_type(data_type):
+ raise ValueError("Policy type {} is not written as Java writes
it.".format(json.dumps(value)))
+ return data_type
+
+
+def _declared_types(transform: Dict[str, Any]) -> Dict[str, DataType]:
+ return {ref["name"]: _declared_type(ref["type"]) for ref in
_field_refs(transform)}
+
+
+def _canonical_transform(transform: Dict[str, Any], schema: TableSchema) ->
Dict[str, Any]:
+ names = [field.name for field in schema.fields]
+
+ def remap(value: Any) -> Any:
+ if not isinstance(value, dict):
+ return value
+ index = names.index(value["name"])
+ return {"index": index, "name": value["name"], "type":
_java_type(schema.fields[index].type)}
+
+ canonical = {"name": transform["name"]}
+ if "fieldRef" in transform:
+ canonical["fieldRef"] = remap(transform["fieldRef"])
+ if transform["name"] == "NESTED_FIELD_REF":
+ canonical["path"] = transform["path"]
+ if transform["name"] == "CAST":
+ canonical["type"] = _java_type(_declared_type(transform["type"]))
+ elif transform["name"] != "NULL":
+ canonical["inputs"] = [remap(value) for value in
transform.get("inputs") or []]
+ if transform["name"] in _TRANSFORM_FLAGS:
+ flag = _TRANSFORM_FLAGS[transform["name"]][0]
+ canonical[flag] = transform[flag]
+ return canonical
+
+
+def _canonical_predicate(predicate: Dict[str, Any], schema: TableSchema) ->
Dict[str, Any]:
+ if predicate["kind"] == "COMPOUND":
+ children = [_canonical_predicate(child, schema) for child in
predicate["children"]]
+ return children[0] if len(children) == 1 else {
+ "kind": "COMPOUND", "function": predicate["function"], "children":
children}
+ transform = predicate["transform"]
+ literals = _leaf_literals(predicate, _transform_type(transform))
+ output = _transform_type(transform, {field.name: field.type for field in
schema.fields})
+ literal_type = _literal_type(predicate["function"], output, RuntimeError)
+ return {"kind": "LEAF", "transform": _canonical_transform(transform,
schema), "function": predicate["function"],
+ "literals": None if literals is None else
[_json_literal(literal_type, value) for value in literals]}
+
+
+def _canonical_policy(policy: DataPolicy, schema: TableSchema) -> DataPolicy:
+ mask = policy.get_column_mask()
+ if mask is None:
+ predicate =
_canonical_predicate(_policy_json(policy.get_row_filter().get_predicate()),
schema)
+ return DataPolicy.row_filter(policy.get_resource(),
RowFilter(_flat_json(predicate)), policy.get_principal())
+ transform = _canonical_transform(_policy_json(mask.get_transform()),
schema)
+ return DataPolicy.column_mask(policy.get_resource(),
ColumnMask(mask.get_on_column(), _flat_json(transform)),
+ policy.get_principal())
+
+
+def _flat_json(value: Any) -> str:
+ if isinstance(value, _Binary):
+ data = value.text.encode("utf-8")
+ encoded = base64.b64encode(data).decode("ascii")
+ value = {"segments": [{"heapMemory": encoded, "array": encoded,
"offHeap": False}],
+ "offset": 0, "sizeInBytes": len(data)}
+ elif isinstance(value, _Decimal):
+ value = {"compact": value.precision <= 18}
+ elif isinstance(value, _Timestamp):
+ value = {"millisecond": value.millisecond, "nanoOfMillisecond":
value.nano_of_millisecond}
+ if isinstance(value, dict):
+ return "{" + ",".join(_flat_json(key) + ":" + _flat_json(item) for
key, item in value.items()) + "}"
+ if isinstance(value, list):
+ return "[" + ",".join(map(_flat_json, value)) + "]"
+ if isinstance(value, float):
+ text = _java_number_string(value, isinstance(value, _JavaFloat))
+ return json.dumps(text) if math.isnan(value) or math.isinf(value) else
text
+ return json.dumps(value, ensure_ascii=False)
+
+
+_JAVA_TYPE_FORMATS = {
+ casts.CHAR: "CHAR({})", casts.VARCHAR: "VARCHAR({})", casts.BINARY:
"BINARY({})",
+ casts.VARBINARY: "VARBINARY({})", casts.INTEGER: "INT", casts.DECIMAL:
"DECIMAL({}, {})",
+ casts.TIME: "TIME({})", casts.TIMESTAMP: "TIMESTAMP({})",
casts.TIMESTAMP_LTZ: "TIMESTAMP({}) WITH LOCAL TIME ZONE"}
+
+
+def _java_type(data_type: DataType) -> Any:
+ suffix = "" if data_type.nullable else " NOT NULL"
+ if isinstance(data_type, AtomicType):
+ root, parameters, _ = _logical_type(data_type)
+ if parameters == (2147483647,) and root in (casts.VARCHAR,
casts.VARBINARY):
+ return ("STRING" if root == casts.VARCHAR else "BYTES") + suffix
+ return _JAVA_TYPE_FORMATS.get(root, root).format(*parameters) + suffix
+ if isinstance(data_type, RowType):
+ return {"type": "ROW" + suffix, "fields": [_java_field(field) for
field in data_type.fields]}
+ if isinstance(data_type, MapType):
+ return {"type": "MAP" + suffix, "key": _java_type(data_type.key),
"value": _java_type(data_type.value)}
+ name = {ArrayType: "ARRAY", MultisetType: "MULTISET", VectorType:
"VECTOR"}[type(data_type)]
+ java = {"type": name + suffix, "element": _java_type(data_type.element)}
+ if isinstance(data_type, VectorType):
+ java["length"] = data_type.length
+ return java
+
+
+def _java_field(field: DataField) -> Dict[str, Any]:
+ java = {"id": field.id, "name": field.name, "type": _java_type(field.type)}
+ if field.description is not None:
+ java["description"] = field.description
+ if field.default_value is not None:
+ java["defaultValue"] = field.default_value
+ return java
+
+
+def _path_ids(data_type: Optional[DataType], path: List[str]) ->
Optional[List[int]]:
+ ids = []
+ for part in path:
+ field = next((field for field in getattr(data_type, "fields", ()) if
field.name == part), None)
+ if field is None:
+ return None
+ ids.append(field.id)
+ data_type = field.type
+ return ids
+
+
+def _moved_columns(policy: DataPolicy, previous: TableSchema, current:
TableSchema) -> List[str]:
+ previous_ids = {field.name: field.id for field in previous.fields}
+ current_ids = {field.name: field.id for field in current.fields}
+ return sorted(column for column in _policy_columns(policy) if
current_ids.get(column) != previous_ids.get(column))
+
+
def _granted_columns(assignment: PermissionAssignment) -> List[str]:
columns = assignment.get_columns()
names = columns.get_column_names()
@@ -267,6 +875,14 @@ def _matches(parameters: Dict[str, str], name: str, value:
Optional[str]) -> boo
return name not in parameters or parameters[name] == value
+def _matches_policy(policy: DataPolicy, parameters: Dict[str, str]) -> bool:
+ mask = policy.get_column_mask()
+ if "column" in parameters and (mask is None or mask.get_on_column() !=
parameters["column"]):
+ return False
+ return _matches(parameters, "type", policy.type().name) \
+ and _matches(parameters, "principal", policy.get_principal())
+
+
class RESTCatalogServer:
"""Mock REST server for testing"""
@@ -296,7 +912,11 @@ class RESTCatalogServer:
self.table_token_store: Dict[str, "RESTToken"] = {}
self.table_query_auth_store = {}
self.permission_store: Dict[str, PermissionAssignment] = {}
+ self.policy_store: Dict[str, Dict[str, DataPolicy]] = {}
+ self.table_identifiers: Dict[str, Identifier] = {}
self.query_principals: set = set()
+ self.management_principals: set = set()
+ self.no_management_permission_resources: set = set()
# Initialize mock catalog (simplified)
self.data_path = data_path
@@ -490,6 +1110,12 @@ class RESTCatalogServer:
path_parts = resource_path[len(self.database_uri) +
1:].split('/')
database_name = RESTUtil.decode_string(path_parts[0])
+ if len(path_parts) in (4, 5) and path_parts[1] ==
ResourcePaths.TABLES \
+ and path_parts[3] == ResourcePaths.POLICIES and
path_parts[4:] in ([], ['drop']):
+ return self._policies_handle(
+ method, data, Identifier.create(database_name,
RESTUtil.decode_string(path_parts[2])),
+ parameters, len(path_parts) == 5)
+
# Check database permissions
if database_name in self.no_permission_databases:
raise DatabaseNoPermissionException(database_name)
@@ -638,6 +1264,8 @@ class RESTCatalogServer:
elif operation == "token":
return self._table_token_handle(method, lookup_identifier)
elif operation == "auth":
+ if branch_part is not None:
+ self._branch_table_metadata(lookup_identifier, branch_part)
return self._table_query_auth_handle(method, data,
lookup_identifier)
elif operation == "rollback":
return self._table_rollback_handle(method, data,
lookup_identifier)
@@ -669,15 +1297,35 @@ class RESTCatalogServer:
def set_query_principals(self, *principals: str) -> None:
self.query_principals = set(principals)
+ def register_management_principal(self, *principals: str) -> None:
+ self.management_principals.update(principals)
+
+ def deny_management_permission(self, resource: PermissionResource) -> None:
+ self.no_management_permission_resources.add(resource)
+
+ def _denied_management(self, resource: PermissionResource) ->
Optional[Tuple[str, int]]:
+ if resource not in self.no_management_permission_resources:
+ return None
+ return self._mock_response(ErrorResponse(
+ resource.get_type().name, _resource_name(resource),
+ "The caller cannot manage permissions on this resource.", 403),
403)
+
+ def _missing_principal(self, principal: str) -> Optional[Tuple[str, int]]:
+ if principal in self.management_principals:
+ return None
+ return self._mock_response(ErrorResponse(
+ "PRINCIPAL", principal, "Permission principal does not exist.",
404), 404)
+
def _permissions_handle(self, method: str, resource_path: str, data: str,
parameters: Dict[str, str]) -> Tuple[str, int]:
if method == "GET" and resource_path ==
self.resource_paths.permissions():
target = PermissionResource(
parameters.get("resourceType"), parameters.get("database"),
parameters.get("table"), parameters.get("function"),
parameters.get("view"))
- missing = self._missing_permission_resource(target)
- if missing is not None:
- return missing
+ error = self._denied_management(target) or
self._missing_permission_resource(target) \
+ or (self._missing_principal(parameters["principal"]) if
"principal" in parameters else None)
+ if error is not None:
+ return error
matched = [assignment for assignment in
self.permission_store.values()
if assignment.get_resource() == target
and _matches(parameters, "principal",
assignment.get_principal())
@@ -686,19 +1334,20 @@ class RESTCatalogServer:
*self._management_page({_assignment_key(grant): grant for
grant in matched}, parameters)), 200)
if method == "POST" and resource_path ==
self.resource_paths.grant_permission():
assignment = JSON.from_json(data,
GrantPermissionRequest).assignment()
- missing =
self._missing_permission_resource(assignment.get_resource())
- if missing is not None:
- return missing
- invalid = self._invalid_column_grant(assignment)
- if invalid is not None:
- return invalid
+ error = self._denied_management(assignment.get_resource()) \
+ or
self._missing_permission_resource(assignment.get_resource()) \
+ or self._missing_principal(assignment.get_principal()) or
self._invalid_column_grant(assignment)
+ if error is not None:
+ return error
self.permission_store[_assignment_key(assignment)] = assignment
return self._mock_response("", 200)
if method == "POST" and resource_path ==
self.resource_paths.revoke_permission():
request = JSON.from_json(data, RevokePermissionRequest)
- missing = self._missing_permission_resource(request.get_resource())
- if missing is not None:
- return missing
+ error = self._denied_management(request.get_resource()) \
+ or self._missing_permission_resource(request.get_resource()) \
+ or self._missing_principal(request.get_principal())
+ if error is not None:
+ return error
self.permission_store.pop(_permission_key(
request.get_resource(), request.get_access(),
request.get_principal()), None)
return self._mock_response("", 200)
@@ -747,26 +1396,102 @@ class RESTCatalogServer:
"Permission column does not exist.", 404), 404)
return None
- def _restricted(self, identifier: Identifier) -> bool:
- return bool(self._column_grants(identifier))
+ def _restricted(self, identifier: Identifier, metadata: TableMetadata) ->
bool:
+ return bool(self.policy_store.get(metadata.uuid or
identifier.get_full_name())
+ or self._column_grants(identifier))
def _column_grants(self, identifier: Identifier) ->
List[PermissionAssignment]:
return [assignment for assignment in self.permission_store.values()
if assignment.get_resource().get_type() == ResourceType.COLUMN
and _on_table(assignment.get_resource(), identifier)]
+ def _policies_handle(self, method: str, data: str, identifier: Identifier,
+ parameters: Dict[str, str], drop: bool) -> Tuple[str,
int]:
+ resource = PermissionResource(
+ ResourceType.TABLE, identifier.get_database_name(),
identifier.get_object_name())
+ denied = self._denied_management(resource)
+ if denied is not None:
+ return denied
+ metadata = self.table_metadata_store.get(identifier.get_full_name())
+ if metadata is None:
+ return self._missing_permission_resource(resource)
+ policies = self._table_policies(identifier, metadata)
+ if method == "GET" and not drop:
+ unknown = self._missing_principal(parameters["principal"]) if
"principal" in parameters else None
+ if unknown is not None:
+ return unknown
+ matched = [DataPolicy(resource, policy.get_row_filter(),
policy.get_column_mask(),
+ policy.get_principal())
+ for policy in policies.values() if
_matches_policy(policy, parameters)]
+ return self._mock_response(ListPoliciesResponse(
+ *self._management_page({_policy_key(policy): policy for policy
in matched}, parameters)), 200)
+ if method == "POST" and not drop:
+ policy = JSON.from_json(data, PolicyRequest).policy(resource)
+ if self.branch_store.get(identifier.get_full_name()):
+ return self._mock_response(ErrorResponse(
+ ErrorResponse.RESOURCE_TYPE_TABLE,
identifier.get_full_name(),
+ "Data policies are not supported on tables with
branches.", 400), 400)
+ if not _query_auth_enabled(metadata.schema):
+ return self._mock_response(ErrorResponse(
+ ErrorResponse.RESOURCE_TYPE_TABLE,
identifier.get_full_name(),
+ "Data policies require the target table option
query-auth.enabled=true.",
+ 409), 409)
+ unknown = self._missing_principal(policy.get_principal())
+ if unknown is not None:
+ return unknown
+ identity = _policy_identity(policy)
+ try:
+ missing = sorted(_policy_columns(policy) - {field.name for
field in metadata.schema.fields})
+ error = "Policy column {} does not exist in table {}.".format(
+ missing[0], identifier.get_full_name()) if missing else
_policy_type_error(policy, metadata.schema)
+ except Exception as e:
+ error = "Invalid policy: {}".format(e)
+ if error:
+ return self._mock_response(ErrorResponse(
+ ErrorResponse.RESOURCE_TYPE_POLICY,
_policy_name(*identity), error, 400), 400)
+ try:
+ canonical = _canonical_policy(policy, metadata.schema)
+ except _UnsupportedConversion as e:
+ return self._mock_response(ErrorResponse(
+ ErrorResponse.RESOURCE_TYPE_POLICY,
_policy_name(*identity), str(e), 400), 400)
+ if _cursor_key(*identity) in policies:
+ return self._mock_response(ErrorResponse(
+ ErrorResponse.RESOURCE_TYPE_POLICY,
_policy_name(*identity),
+ "Policy already exists.", 409), 409)
+ policies[_cursor_key(*identity)] = canonical
+ return self._mock_response("", 200)
+ if method == "POST" and drop:
+ request = JSON.from_json(data, DropPolicyRequest)
+ identity = (request.get_type().name, request.get_principal(),
request.get_column())
+ if policies.pop(_cursor_key(*identity), None) is None:
+ return self._mock_response(ErrorResponse(
+ ErrorResponse.RESOURCE_TYPE_POLICY,
_policy_name(*identity),
+ "Policy does not exist.", 404), 404)
+ return self._mock_response("", 200)
+ return self._mock_response(ErrorResponse(None, None, "Not Found",
404), 404)
+
def _remove_permissions(self, matches) -> None:
for key in [key for key, assignment in self.permission_store.items()
if matches(assignment.get_resource())]:
del self.permission_store[key]
- def _drop_table_access(self, identifier: Identifier) -> None:
+ def _drop_table_access(self, identifier: Identifier, metadata:
TableMetadata) -> None:
for branch in self.branch_store.pop(identifier.get_full_name(), ()):
manager = self._schema_manager(identifier, branch)
manager.file_io.delete(manager.branch_path, recursive=True)
+ if not metadata.is_external:
+ manager = self._schema_manager(identifier)
+ manager.file_io.delete(manager.table_path, recursive=True)
+ self.policy_store.pop(metadata.uuid or identifier.get_full_name(),
None)
+ self.table_identifiers.pop(identifier.get_full_name(), None)
self._remove_permissions(lambda resource: _on_table(resource,
identifier))
+ def _table_identifier(self, name: str) -> Identifier:
+ return self.table_identifiers.get(name) or Identifier.from_string(name)
+
def _rename_table_access(self, source: Identifier, destination:
Identifier) -> None:
+ self.table_identifiers.pop(source.get_full_name(), None)
+ self.table_identifiers[destination.get_full_name()] = destination
moved = [assignment for assignment in self.permission_store.values()
if _on_table(assignment.get_resource(), source)]
self._remove_permissions(lambda resource: _on_table(resource, source))
@@ -777,6 +1502,15 @@ class RESTCatalogServer:
assignment.get_access(), assignment.get_principal(),
assignment.get_columns(), assignment.get_expire_time())
self.permission_store[_assignment_key(renamed)] = renamed
+ metadata = self.table_metadata_store[destination.get_full_name()]
+ source_key = metadata.uuid or source.get_full_name()
+ if source_key in self.policy_store:
+ policies = self.policy_store.pop(source_key)
+ self._table_policies(destination, metadata).update(policies)
+
+ def _table_policies(self, identifier: Identifier,
+ metadata: TableMetadata) -> Dict[str, DataPolicy]:
+ return self.policy_store.setdefault(metadata.uuid or
identifier.get_full_name(), {})
def _check_column_grants(self, identifier: Identifier, select:
Optional[List[str]]) -> None:
fields = [field.name for field in
self.table_metadata_store[identifier.get_full_name()].schema.fields]
@@ -793,6 +1527,29 @@ class RESTCatalogServer:
if any(column not in included for column in (fields if select is
None else select)):
raise TableNoPermissionException(identifier)
+ def _policy_rules(self, identifier: Identifier) ->
Tuple[Optional[List[str]], Optional[Dict[str, str]]]:
+ metadata = self.table_metadata_store[identifier.get_full_name()]
+ policies = self._table_policies(identifier, metadata)
+ names = {field.name for field in metadata.schema.fields}
+ applied, masked = [], set()
+ for key in sorted(policies):
+ policy = policies[key]
+ if policy.get_principal() not in self.query_principals:
+ continue
+ try:
+ invalid = _policy_columns(policy) - names or
_policy_type_error(policy, metadata.schema)
+ except Exception:
+ invalid = True
+ mask = policy.get_column_mask()
+ if invalid or (mask is not None and mask.get_on_column() in
masked):
+ raise TableNoPermissionException(identifier)
+ masked |= {mask.get_on_column()} if mask is not None else set()
+ applied.append(_canonical_policy(policy, metadata.schema))
+ filters = [policy.get_row_filter().get_predicate() for policy in
applied if policy.get_row_filter()]
+ masks = {policy.get_column_mask().get_on_column():
policy.get_column_mask().get_transform()
+ for policy in applied if policy.get_column_mask()}
+ return filters or None, masks or None
+
# ======================= Function Handlers ===============================
def _functions_handle(self, method: str, data: str, database_name: str,
@@ -1124,8 +1881,8 @@ class RESTCatalogServer:
raise TagNotExistException(request.from_tag)
snapshot = tags[request.from_tag].snapshot
schema_id = schema_id if snapshot is None else
snapshot.schema_id
- if self._restricted(identifier):
- raise ValueError("Tables with column permissions cannot have
branches.")
+ if self._restricted(identifier, metadata):
+ raise ValueError("Tables with data policies or column
permissions cannot have branches.")
manager = self._schema_manager(identifier)
if snapshot is not None:
self._write_snapshot_files(identifier, snapshot, None,
request.branch)
@@ -1219,6 +1976,11 @@ class RESTCatalogServer:
elif method == "DELETE":
del self.database_store[database_name]
+ for name in [name for name in self.table_metadata_store
+ if self._table_identifier(name).get_database_name()
== database_name]:
+ self._drop_table_access(self._table_identifier(name),
self.table_metadata_store.pop(name))
+ self.table_latest_snapshot_store.pop(name, None)
+ self.table_partitions_store.pop(name, None)
self._remove_permissions(lambda resource: resource.get_database()
== database_name)
return self._mock_response("", 200)
return self._mock_response(ErrorResponse(None, None, "Method Not
Allowed", 405), 405)
@@ -1242,6 +2004,7 @@ class RESTCatalogServer:
create_table.identifier, 0, create_table.schema,
str(uuid.uuid4()), False
)
self.table_metadata_store.update({create_table.identifier.get_full_name():
table_metadata})
+
self.table_identifiers[create_table.identifier.get_full_name()] =
create_table.identifier
table_dir = (
Path(self.data_path) / self.warehouse / database_name /
create_table.identifier.get_object_name() / 'schema'
@@ -1274,11 +2037,12 @@ class RESTCatalogServer:
return self._mock_response("", 200)
elif method == "DELETE":
# Drop table
- if identifier.get_full_name() not in self.table_metadata_store:
- raise TableNotExistException(identifier)
+ if identifier.get_full_name() not in self.table_metadata_store \
+ or (response_identifier or identifier).get_branch_name()
is not None:
+ raise TableNotExistException(response_identifier or identifier)
else:
- del self.table_metadata_store[identifier.get_full_name()]
- self._drop_table_access(identifier)
+ self._drop_table_access(
+ identifier,
self.table_metadata_store.pop(identifier.get_full_name()))
if identifier.get_full_name() in self.table_latest_snapshot_store:
del
self.table_latest_snapshot_store[identifier.get_full_name()]
if identifier.get_full_name() in self.table_partitions_store:
@@ -1296,8 +2060,10 @@ class RESTCatalogServer:
select = json.loads(data).get('select')
self._check_column_grants(identifier, select)
from pypaimon.api.api_response import AuthTableQueryResponse
- response = self.table_query_auth_store.get(identifier.get_full_name(),
AuthTableQueryResponse())
- if callable(response):
+ response = self.table_query_auth_store.get(identifier.get_full_name())
+ if response is None:
+ response = AuthTableQueryResponse(*self._policy_rules(identifier))
+ elif callable(response):
response = response(select)
return self._mock_response(response, 200)
@@ -1633,7 +2399,8 @@ class RESTCatalogServer:
parameters: Dict[str, str]) -> Tuple[List[Any],
Optional[str]]:
token = parameters.get(PAGE_TOKEN)
keys = sorted(key for key in entities if not token or key > token)
- page = keys[:self._get_max_results(parameters)]
+ size = int(parameters[MAX_RESULTS]) if parameters.get(MAX_RESULTS)
else DEFAULT_MAX_RESULTS
+ page = keys[:max(1, min(size, ListPermissionsRequest.MAX_PAGE_SIZE))]
return [entities[key] for key in page], page[-1] if len(keys) >
len(page) else None
def _build_paged_entities(self, entities: List[Any], max_results: int,
@@ -1753,6 +2520,21 @@ class RESTCatalogServer:
def _validate_access_for_schema(self, identifier: Identifier, metadata:
TableMetadata,
schema: TableSchema) -> None:
name = identifier.get_full_name()
+ policies = self.policy_store.get(metadata.uuid or name, {})
+ if policies and not _query_auth_enabled(schema):
+ raise ValueError(
+ "Cannot disable query-auth.enabled while table {} has data
policies.".format(name))
+ for policy in policies.values():
+ missing = _moved_columns(policy, metadata.schema, schema)
+ if missing:
+ raise ValueError("Cannot remove or rename policy column {}
from table {}.".format(missing[0], name))
+ mismatch = _policy_type_error(policy, schema)
+ if mismatch:
+ raise ValueError(mismatch)
+ try:
+ _canonical_policy(policy, schema)
+ except Exception as e:
+ raise ValueError(str(e)) from e
grants = self._column_grants(identifier)
if grants and not _query_auth_enabled(schema):
raise ValueError(
@@ -1826,7 +2608,7 @@ class RESTCatalogServer:
tables = []
for full_name, metadata in self.table_metadata_store.items():
- identifier = Identifier.from_string(full_name)
+ identifier = self._table_identifier(full_name)
metadata_table_type = (
metadata.schema.options.get(TYPE, "table")
if metadata and metadata.schema and metadata.schema.options
diff --git a/paimon-python/pypaimon/tests/sql_permission_test.py
b/paimon-python/pypaimon/tests/sql_permission_test.py
index 158f17f72a..ae5655bdcc 100644
--- a/paimon-python/pypaimon/tests/sql_permission_test.py
+++ b/paimon-python/pypaimon/tests/sql_permission_test.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
+import json
+
import pyarrow as pa
import pytest
@@ -22,11 +24,17 @@ from pypaimon import CatalogFactory, Schema
from pypaimon.api.api_response import ConfigResponse
from pypaimon.api.auth import BearTokenAuthProvider
from pypaimon.catalog.catalog_exception import TableNoPermissionException
+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.policy_management import PolicyAlreadyExistException
+from pypaimon.management.policy_type import PolicyType
from pypaimon.management.resource_type import ResourceType
+from pypaimon.management.row_filter import RowFilter
from pypaimon.read.native_plan import native_runtime_available
from pypaimon.schema.schema_change import SchemaChange
from pypaimon.tests.rest.rest_server import RESTCatalogServer
@@ -46,10 +54,19 @@ pytestmark = [pytest.mark.native_plan, pytest.mark.skipif(
TOKEN = 'test-token'
ANALYST = 'analyst'
ROWS = [
- {'id': 1, 'region': 'eu', 'email': '[email protected]'},
- {'id': 2, 'region': 'us', 'email': '[email protected]'},
- {'id': 3, 'region': 'eu', 'email': '[email protected]'},
+ {'id': 1, 'region': 'eu', 'email': '[email protected]', 'vip': True},
+ {'id': 2, 'region': 'us', 'email': '[email protected]', 'vip': False},
+ {'id': 3, 'region': 'eu', 'email': '[email protected]', 'vip': None},
]
+EU_ONLY = json.dumps({
+ 'kind': 'LEAF',
+ 'transform': {'name': 'FIELD_REF',
+ 'fieldRef': {'index': 1, 'name': 'region', 'type':
'STRING'}},
+ 'function': 'EQUAL',
+ 'literals': ['eu'],
+}, separators=(',', ':'))
+NULL_MASK = json.dumps({'name': 'NULL'}, separators=(',', ':'))
+STARS_MASK = json.dumps({'name': 'CONCAT', 'inputs': ['***']},
separators=(',', ':'))
ORDERS = PermissionResource(ResourceType.TABLE, 'default', 'orders')
@@ -60,6 +77,7 @@ def clients(tmp_path):
server = RESTCatalogServer(str(tmp_path), BearTokenAuthProvider(TOKEN),
ConfigResponse(defaults={'prefix':
'permission-test'}),
'warehouse')
+ server.register_management_principal(ANALYST)
server.start()
try:
options = {'metastore': 'rest', 'uri': server.get_url(), 'warehouse':
'warehouse',
@@ -74,7 +92,7 @@ def clients(tmp_path):
def _create_orders(catalog, rows=()):
- schema = pa.schema([('id', pa.int64()), ('region', pa.string()), ('email',
pa.string())])
+ schema = pa.schema([('id', pa.int64()), ('region', pa.string()), ('email',
pa.string()), ('vip', pa.bool_())])
catalog.create_table('default.orders', Schema.from_pyarrow_schema(schema),
False)
if rows:
table = catalog.get_table('default.orders')
@@ -97,14 +115,14 @@ def _call_ok(context, sql):
assert _sql(context, sql) == [{'result': 'OK'}], sql
-def _python_rows(catalog, columns=('id', 'region', 'email')):
+def _python_rows(catalog, columns=('id', 'region', 'email', 'vip')):
table = catalog.get_table('default.orders')
read_builder = table.new_read_builder().with_projection(list(columns))
splits = read_builder.new_scan().plan().splits()
return sorted(read_builder.new_read().to_arrow(splits).to_pylist(),
key=lambda row: row['id'])
-def _sql_rows(context, columns=('id', 'region', 'email')):
+def _sql_rows(context, columns=('id', 'region', 'email', 'vip')):
return _sql(context, 'SELECT {} FROM paimon.default.orders ORDER BY
id'.format(', '.join(columns)))
@@ -159,3 +177,83 @@ def test_column_grants_restrict_both_readers(clients):
_python_rows(catalog)
with pytest.raises(Exception, match='(?i)permission'):
_sql_rows(context)
+
+
+def _policies(catalog):
+ page =
catalog.policy_management().list_policies(ListPoliciesRequest(ORDERS))
+ return sorted((policy.type().name, policy.get_principal(),
+ policy.get_row_filter() and
policy.get_row_filter().get_predicate(),
+ policy.get_column_mask() and
policy.get_column_mask().get_on_column(),
+ policy.get_column_mask() and
policy.get_column_mask().get_transform())
+ for policy in page.elements)
+
+
+def test_sql_and_pypaimon_manage_the_same_policies(clients):
+ catalog, context, _ = clients
+ _create_orders(catalog)
+ policies = catalog.policy_management()
+
+ _call_ok(context, "CALL sys.create_policy(database => 'default', table =>
'orders', "
+ "policy_type => 'ROW_FILTER', principal => 'analyst', "
+ "predicate_json => '{}')".format(EU_ONLY))
+ policies.create_policy(DataPolicy.column_mask(ORDERS, ColumnMask('vip',
NULL_MASK), ANALYST))
+ assert _policies(catalog) == [
+ ('COLUMN_MASKING', ANALYST, None, 'vip', NULL_MASK),
+ ('ROW_FILTER', ANALYST, EU_ONLY, None, None),
+ ]
+ assert sorted((row['policy_type'], row['on_column'],
row['predicate_json'], row['transform_json'])
+ for row in _sql(context, "CALL sys.list_policies(database =>
'default', "
+ "table => 'orders', principal =>
'analyst')")) == [
+ ('COLUMN_MASKING', 'vip', None, NULL_MASK),
+ ('ROW_FILTER', None, EU_ONLY, None),
+ ]
+
+ with pytest.raises(PolicyAlreadyExistException):
+ policies.create_policy(DataPolicy.row_filter(ORDERS,
RowFilter(EU_ONLY), ANALYST))
+ with pytest.raises(Exception, match='already exists'):
+ context.sql("CALL sys.create_policy(database => 'default', table =>
'orders', "
+ "policy_type => 'COLUMN_MASKING', principal => 'analyst', "
+ "on_column => 'vip', transform_json =>
'{}')".format(NULL_MASK))
+
+ _call_ok(context, "CALL sys.drop_policy(database => 'default', table =>
'orders', "
+ "policy_type => 'COLUMN_MASKING', principal =>
'analyst', column => 'vip')")
+ policies.drop_policy(ORDERS, PolicyType.ROW_FILTER, ANALYST, None, False)
+ assert _policies(catalog) == []
+ _call_ok(context, "CALL sys.drop_policy(database => 'default', table =>
'orders', "
+ "policy_type => 'ROW_FILTER', principal => 'analyst',
if_exists => 'true')")
+
+
[email protected](RUST_VERSION < (0, 5), reason='pypaimon-rust 0.5 or later
required')
+def test_row_filter_restricts_both_readers_alike(clients):
+ catalog, context, server = clients
+ _create_orders(catalog, ROWS)
+ _call_ok(context, "CALL sys.create_policy(database => 'default', table =>
'orders', "
+ "policy_type => 'ROW_FILTER', principal => 'analyst', "
+ "predicate_json => '{}')".format(EU_ONLY))
+
+ assert _python_rows(catalog) == ROWS
+ assert _sql_rows(context) == ROWS
+
+ server.set_query_principals(ANALYST)
+ eu_rows = [row for row in ROWS if row['region'] == 'eu']
+ assert _python_rows(catalog) == eu_rows
+ assert _sql_rows(context) == eu_rows
+ assert _sql(context, 'SELECT count(*) AS n FROM paimon.default.orders') ==
[{'n': 2}]
+
+
[email protected](RUST_VERSION < (0, 5), reason='pypaimon-rust 0.5 or later
required')
+def test_column_masks_mask_both_readers_alike(clients):
+ catalog, context, server = clients
+ _create_orders(catalog, ROWS)
+ policies = catalog.policy_management()
+ policies.create_policy(DataPolicy.column_mask(ORDERS, ColumnMask('vip',
NULL_MASK), ANALYST))
+ policies.create_policy(DataPolicy.column_mask(ORDERS, ColumnMask('region',
STARS_MASK), ANALYST))
+
+ server.set_query_principals(ANALYST)
+ masked = [dict(row, region='***', vip=None) for row in ROWS]
+ assert _python_rows(catalog) == masked
+ assert _sql_rows(context) == masked
+
+ server.set_query_principals('someone-else')
+ assert _python_rows(catalog) == ROWS
+ assert _sql_rows(context) == ROWS
