Add SystemAccessControl.getTableColumnMasks SPI function and OPA implementation #21997

mosiac1 · 2024-05-16T11:05:15Z

Description

Addresses issues raised in #21359.

Adds the getTableColumnMasks function to the SPI

Map<ColumnSchema, ViewExpression> getTableColumnMasks(ConnectorSecurityContext context, SchemaTableName tableName, List<ColumnSchema> columns);

(which defaults to calling getColumnMask repeatedly) that SystemAcessControl/ConnectorAccessControl plugins can implement to process column mask requests in batches.

Updates the StatementAnalyzer and AccessControl to use this batched method of getting column masks whenever possible (which is all cases).

The issue linked above goes into the performance implications.

This PR updates the OPA plugin to implement the new SPI function with 2 modes of operation:

Sending requests to fetch each column mask in parallel; This is a drop-in improvement that doesn't require any config/Rego changes.
Using a new batchColumnMasksUri OPA Server endpoint that delegates collection of column masks to the Rego policy code; This requires changes to the Rego code running on the OPA server and the Trino configs.

For the latter mode, the following Rego rule can be added (assuming the columnMask rule is already implemented) to quickly add support for batched column masks. This approach isn't particularly efficient (the with operator in Rego can be costly) and I haven't had the time to benchmark.
batchColumnMasks contains {
   "index": i,
   "viewExpression": cm
} if {
  some i
  cm := columnMask with input.action.resource.column as input.action.filterResources[i].
}

Later edit: did some testing on using the above "trick" - it is very slow (column masks for 1k columns took on avg. 16 seconds to compute). So DO NOT USE THIS IN PRODUCTION. The parallel mode (which this PR enables by default) is much faster.

Additional context and related issues

#21359

Release notes

( ) This is not user-visible or is docs only, and no release notes are required.
(x) Release notes are required. Please propose a release note for me.
( ) Release notes are required, with the following suggested text:

# Section
* Fix some things. ({issue}`issuenumber`)

mosiac1 · 2024-05-16T13:28:30Z

core/trino-spi/src/main/java/io/trino/spi/security/SystemAccessControl.java

+        return columns.stream()
+                .map(c -> Map.entry(c, getColumnMask(context, tableName, c.getName(), c.getType())))
+                .filter(entry -> entry.getValue().isPresent())
+                .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().get()));


I'm not sure if this should have a default that calls getColumnMask, looking at other function in this interface none of them do this. Should this logic be left to another class?

yes, the default should be to call the old method
the engine should call the new method only
the old method should be deprecated

thanks! do we want to remove getColumnMask from the AccessControl interface? that would be the best way to ensure the engine is not using it anymore

yes, remove the supposed-to-be-unused method from the interface defined in trino-main

deprecation is needed only for trino-spi's interfaces

Removed the function from trino-main's interfaces, let me know if missed some.

Do we want to remove getColumnMask's implementation from the FileBasedAccessControl and OpaAccessControl implementation? (or others if you have any in mind)

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorAccessControl.java

findepi · 2024-05-17T21:15:37Z

core/trino-spi/src/main/java/io/trino/spi/security/SystemAccessControl.java

+        return columns.stream()
+                .map(c -> Map.entry(c, getColumnMask(context, tableName, c.getName(), c.getType())))
+                .filter(entry -> entry.getValue().isPresent())
+                .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().get()));


yes, the default should be to call the old method
the engine should call the new method only
the old method should be deprecated

dain

Looks good.

Let's cleanup the commits as follows:

Squash Deprecate getColumnMask from SPI interfaces into Add access control SPI function for table column masks. The deprecation should be mentioned in the commit message body.
Fixup Remove getColumnMask from the AccessControl interface into Add use table column masks SPI function for access control in core and change title to say something like change access control to fetch column masks in bulk
Fixup Add FBAC tests for getTableColumnMasks into the commit that changed file based access control (no need to mention you added tests in the commit message/body)
For the other two commits change the commit message to say something like "Implement column mask bulk fetch api" or "convert to ...", so it is clear what changed

dain · 2024-06-03T21:46:03Z

core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java

+                field.getName().ifPresent(
+                        fieldName -> columnSchemaBuilder.add(ColumnSchema.builder().setName(fieldName).setType(field.getType()).setHidden(field.isHidden()).build()));


I would reformat this so the columnschema builder is split across multiple lines, like this:

Suggested change

field.getName().ifPresent(

fieldName -> columnSchemaBuilder.add(ColumnSchema.builder().setName(fieldName).setType(field.getType()).setHidden(field.isHidden()).build()));

field.getName().ifPresent(fieldName -> columnSchemaBuilder.add(ColumnSchema.builder()

.setName(fieldName)

.setType(field.getType())

.setHidden(field.isHidden())

.build()));

dain · 2024-06-03T21:48:52Z

core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java

-                        analyzeColumnMask(session.getIdentity().getUser(), table, name, field, accessControlScope, mask.get());
-                    }
+            Map<ColumnSchema, ViewExpression> masks = accessControl.getTableColumnMasks(session.toSecurityContext(), name, columnSchemas);
+            for (Map.Entry<ColumnSchema, ViewExpression> entry : masks.entrySet()) {


I would iterate over columnSchemas instead of the returned map incase connector returns extra junk columns... just a bit safer

dain · 2024-06-03T21:49:59Z

plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaAccessControl.java

@@ -739,9 +741,17 @@ public List<ViewExpression> getRowFilters(SystemSecurityContext context, Catalog
    @Override
    public Optional<ViewExpression> getColumnMask(SystemSecurityContext context, CatalogSchemaTableName tableName, String columnName, Type type)


You should be able to drop this method

dain · 2024-06-03T21:51:47Z

plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java

@@ -125,4 +126,18 @@ public OpaConfig setOpaColumnMaskingUri(URI opaColumnMaskingUri)
        this.opaColumnMaskingUri = Optional.ofNullable(opaColumnMaskingUri);
        return this;
    }
+
+    @NotNull
+    public Optional<URI> getOpaBatchColumnMaskingUri()


Are you planning on supporting non-batched column masks in the long run. With this PR Trino will always ask for a batch of masks, do you want to support doing a call per column?

We won't be supporting single column masks. I do think this should be Optional so we don't break current configurations; the plugin will query opaColumnMaskingUri in parallel to get each mask. Admins can set this once they implement batching in Rego.

dain · 2024-06-03T21:55:43Z

plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHighLevelClient.java

+        return result.result().stream()
+                .map(resultItem -> Map.entry(columns.get(resultItem.index()), resultItem.viewExpression()))
+                .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));


This seems a bit redundant:

Suggested change

return result.result().stream()

.map(resultItem -> Map.entry(columns.get(resultItem.index()), resultItem.viewExpression()))

.collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));

return result.result().stream()

.collect(toImmutableMap(item -> columns.get(item.index(), item-> item.viewExpression()));

dain · 2024-06-03T21:57:32Z

plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaBatchColumnMaskQueryResult.java

+    public record OpaBatchColumnMaskQueryResultItem(int index, OpaViewExpression viewExpression)
+    {}


NIT the {} can be on the same line as the record declaration

Deprecate `SystemAccessControl#getColumnMask` and `ConnectorAccessControl#getColumnMask` from SPI interfaces.

cla-bot bot added the cla-signed label May 16, 2024

mosiac1 mentioned this pull request May 16, 2024

When utilizing column masking, unnecessary requests are generated for each column. #21359

Open

mosiac1 commented May 16, 2024

View reviewed changes

findepi reviewed May 17, 2024

View reviewed changes

mosiac1 force-pushed the trinodb/column-masks branch from e256f67 to a31b864 Compare May 21, 2024 12:31

dain reviewed Jun 3, 2024

View reviewed changes

hashhar requested review from kokosing and lukasz-walkiewicz June 4, 2024 12:21

mosiac1 force-pushed the trinodb/column-masks branch from a31b864 to 07a396a Compare June 7, 2024 08:43

github-actions bot added the hive Hive connector label Jun 7, 2024

mosiac1 force-pushed the trinodb/column-masks branch 2 times, most recently from 711ee6a to 9aba72b Compare June 7, 2024 10:02

mosiac1 added 5 commits June 7, 2024 12:32

Add access control SPI function for table column masks

fb5a16c

Deprecate `SystemAccessControl#getColumnMask` and `ConnectorAccessControl#getColumnMask` from SPI interfaces.

Update core access control to fetch column masks in bulk

6fdbf9c

Implement column masks bulk fetch API in trino-opa

9fcfb80

Implement column masks bulk fetch API in trino-plugin-toolkit

d59b30a

Implement column masks bulk fetch API in trino-hive

da4616a

mosiac1 force-pushed the trinodb/column-masks branch from 9aba72b to da4616a Compare June 7, 2024 11:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add SystemAccessControl.getTableColumnMasks SPI function and OPA implementation #21997

Add SystemAccessControl.getTableColumnMasks SPI function and OPA implementation #21997

mosiac1 commented May 16, 2024 •

edited

mosiac1 May 16, 2024

findepi May 17, 2024

mosiac1 May 20, 2024

findepi May 20, 2024

mosiac1 May 21, 2024

findepi May 17, 2024

dain left a comment

dain Jun 3, 2024

dain Jun 3, 2024

dain Jun 3, 2024

dain Jun 3, 2024

mosiac1 Jun 4, 2024

dain Jun 3, 2024

dain Jun 3, 2024

		field.getName().ifPresent(
		fieldName -> columnSchemaBuilder.add(ColumnSchema.builder().setName(fieldName).setType(field.getType()).setHidden(field.isHidden()).build()));

		@@ -739,9 +741,17 @@ public List<ViewExpression> getRowFilters(SystemSecurityContext context, Catalog
		@Override
		public Optional<ViewExpression> getColumnMask(SystemSecurityContext context, CatalogSchemaTableName tableName, String columnName, Type type)

		public record OpaBatchColumnMaskQueryResultItem(int index, OpaViewExpression viewExpression)
		{}

Add SystemAccessControl.getTableColumnMasks SPI function and OPA implementation #21997

Are you sure you want to change the base?

Add SystemAccessControl.getTableColumnMasks SPI function and OPA implementation #21997

Conversation

mosiac1 commented May 16, 2024 • edited

Description

Additional context and related issues

Release notes

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dain left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mosiac1 commented May 16, 2024 •

edited