Implement row level filtering and masking and add system tests

bloomberg · Dec 21, 2023 · 701c0e0 · 701c0e0
1 parent fd6d65c
commit 701c0e0
Show file tree

Hide file tree

Showing 20 changed files with 1,250 additions and 189 deletions.
diff --git a/plugin/trino-opa/pom.xml b/plugin/trino-opa/pom.xml
@@ -121,11 +121,19 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-main</artifactId>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-testing</artifactId>
             <scope>test</scope>
         </dependency>
+
         <dependency>
             <groupId>org.assertj</groupId>
             <artifactId>assertj-core</artifactId>

diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaAccessControl.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaAccessControl.java
@@ -22,6 +22,7 @@
 import io.trino.plugin.opa.schema.OpaQueryInput;
 import io.trino.plugin.opa.schema.OpaQueryInputAction;
 import io.trino.plugin.opa.schema.OpaQueryInputResource;
+import io.trino.plugin.opa.schema.OpaViewExpression;
 import io.trino.plugin.opa.schema.TrinoCatalogSessionProperty;
 import io.trino.plugin.opa.schema.TrinoFunction;
 import io.trino.plugin.opa.schema.TrinoGrantPrincipal;
@@ -40,15 +41,19 @@
 import io.trino.spi.security.SystemAccessControl;
 import io.trino.spi.security.SystemSecurityContext;
 import io.trino.spi.security.TrinoPrincipal;
+import io.trino.spi.security.ViewExpression;
+import io.trino.spi.type.Type;
 
 import java.security.Principal;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static io.trino.plugin.opa.OpaHighLevelClient.buildQueryInputForSimpleResource;
@@ -709,6 +714,23 @@ public void checkCanDropFunction(SystemSecurityContext systemSecurityContext, Ca
                 OpaQueryInputResource.builder().function(TrinoFunction.fromTrinoFunction(functionName)).build());
     }
 
+    @Override
+    public List<ViewExpression> getRowFilters(SystemSecurityContext context, CatalogSchemaTableName tableName)
+    {
+        List<OpaViewExpression> rowFilterExpressions = opaHighLevelClient.getRowFilterExpressionsFromOpa(buildQueryContext(context), tableName);
+        return rowFilterExpressions.stream()
+                .map(expression -> expression.toTrinoViewExpression(tableName.getCatalogName(), tableName.getSchemaTableName().getSchemaName()))
+                .collect(toImmutableList());
+    }
+
+    @Override
+    public Optional<ViewExpression> getColumnMask(SystemSecurityContext context, CatalogSchemaTableName tableName, String columnName, Type type)
+    {
+        return opaHighLevelClient
+                .getColumnMaskFromOpa(buildQueryContext(context), tableName, columnName, type)
+                .map(expression -> expression.toTrinoViewExpression(tableName.getCatalogName(), tableName.getSchemaTableName().getSchemaName()));
+    }
+
     private void checkTableOperation(SystemSecurityContext context, String actionName, CatalogSchemaTableName table, Consumer<String> deny)
     {
         opaHighLevelClient.queryAndEnforce(

diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaAccessControlFactory.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaAccessControlFactory.java
@@ -22,9 +22,11 @@
 import io.airlift.concurrent.BoundedExecutor;
 import io.airlift.http.client.HttpClient;
 import io.airlift.json.JsonModule;
+import io.trino.plugin.opa.schema.OpaColumnMaskQueryResult;
 import io.trino.plugin.opa.schema.OpaPluginContext;
 import io.trino.plugin.opa.schema.OpaQuery;
 import io.trino.plugin.opa.schema.OpaQueryResult;
+import io.trino.plugin.opa.schema.OpaRowFiltersQueryResult;
 import io.trino.spi.security.SystemAccessControl;
 import io.trino.spi.security.SystemAccessControlFactory;
 
@@ -71,6 +73,8 @@ protected static SystemAccessControl create(Map<String, String> config, Optional
                 binder -> {
                     jsonCodecBinder(binder).bindJsonCodec(OpaQuery.class);
                     jsonCodecBinder(binder).bindJsonCodec(OpaQueryResult.class);
+                    jsonCodecBinder(binder).bindJsonCodec(OpaRowFiltersQueryResult.class);
+                    jsonCodecBinder(binder).bindJsonCodec(OpaColumnMaskQueryResult.class);
                     httpClient.ifPresentOrElse(
                             client -> binder.bind(Key.get(HttpClient.class, ForOpa.class)).toInstance(client),
                             () -> httpClientBinder(binder).bindHttpClient("opa", ForOpa.class));

diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java
@@ -28,6 +28,8 @@ public class OpaConfig
     private boolean logRequests;
     private boolean logResponses;
     private boolean allowPermissioningOperations;
+    private Optional<URI> opaRowFiltersUri = Optional.empty();
+    private Optional<URI> opaColumnMaskingUri = Optional.empty();
 
     @NotNull
     public URI getOpaUri()
@@ -43,6 +45,7 @@ public OpaConfig setOpaUri(@NotNull URI opaUri)
         return this;
     }
 
+    @NotNull
     public Optional<URI> getOpaBatchUri()
     {
         return opaBatchUri;
@@ -94,4 +97,32 @@ public OpaConfig setAllowPermissioningOperations(boolean allowPermissioningOpera
         this.allowPermissioningOperations = allowPermissioningOperations;
         return this;
     }
+
+    @NotNull
+    public Optional<URI> getOpaRowFiltersUri()
+    {
+        return opaRowFiltersUri;
+    }
+
+    @Config("opa.policy.row-filters-uri")
+    @ConfigDescription("URI for fetching row filters - if not set no row filtering will be applied")
+    public OpaConfig setOpaRowFiltersUri(@NotNull URI opaRowFiltersUri)
+    {
+        this.opaRowFiltersUri = Optional.ofNullable(opaRowFiltersUri);
+        return this;
+    }
+
+    @NotNull
+    public Optional<URI> getOpaColumnMaskingUri()
+    {
+        return opaColumnMaskingUri;
+    }
+
+    @Config("opa.policy.column-masking-uri")
+    @ConfigDescription("URI for fetching column masks - if not set no masking will be applied")
+    public OpaConfig setOpaColumnMaskingUri(URI opaColumnMaskingUri)
+    {
+        this.opaColumnMaskingUri = Optional.ofNullable(opaColumnMaskingUri);
+        return this;
+    }
 }
diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHighLevelClient.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHighLevelClient.java
@@ -13,17 +13,27 @@
  */
 package io.trino.plugin.opa;
 
+import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 import io.airlift.json.JsonCodec;
+import io.trino.plugin.opa.schema.OpaColumnMaskQueryResult;
 import io.trino.plugin.opa.schema.OpaQueryContext;
 import io.trino.plugin.opa.schema.OpaQueryInput;
 import io.trino.plugin.opa.schema.OpaQueryInputAction;
 import io.trino.plugin.opa.schema.OpaQueryInputResource;
 import io.trino.plugin.opa.schema.OpaQueryResult;
+import io.trino.plugin.opa.schema.OpaRowFiltersQueryResult;
+import io.trino.plugin.opa.schema.OpaViewExpression;
+import io.trino.plugin.opa.schema.TrinoColumn;
+import io.trino.plugin.opa.schema.TrinoTable;
+import io.trino.spi.connector.CatalogSchemaTableName;
 import io.trino.spi.security.AccessDeniedException;
+import io.trino.spi.type.Type;
 
 import java.net.URI;
 import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 
@@ -32,18 +42,28 @@
 public class OpaHighLevelClient
 {
     private final JsonCodec<OpaQueryResult> queryResultCodec;
-    private final URI opaPolicyUri;
+    private final JsonCodec<OpaRowFiltersQueryResult> rowFiltersQueryResultCodec;
+    private final JsonCodec<OpaColumnMaskQueryResult> columnMaskQueryResultCodec;
     private final OpaHttpClient opaHttpClient;
+    private final URI opaPolicyUri;
+    private final Optional<URI> opaRowFiltersUri;
+    private final Optional<URI> opaColumnMaskingUri;
 
     @Inject
     public OpaHighLevelClient(
             JsonCodec<OpaQueryResult> queryResultCodec,
+            JsonCodec<OpaRowFiltersQueryResult> rowFiltersQueryResultCodec,
+            JsonCodec<OpaColumnMaskQueryResult> columnMaskQueryResultCodec,
             OpaHttpClient opaHttpClient,
             OpaConfig config)
     {
         this.queryResultCodec = requireNonNull(queryResultCodec, "queryResultCodec is null");
+        this.rowFiltersQueryResultCodec = requireNonNull(rowFiltersQueryResultCodec, "rowFiltersQueryResultCodec is null");
+        this.columnMaskQueryResultCodec = requireNonNull(columnMaskQueryResultCodec, "columnMaskQueryResultCodec is null");
         this.opaHttpClient = requireNonNull(opaHttpClient, "opaHttpClient is null");
         this.opaPolicyUri = config.getOpaUri();
+        this.opaRowFiltersUri = config.getOpaRowFiltersUri();
+        this.opaColumnMaskingUri = config.getOpaColumnMaskingUri();
     }
 
     public boolean queryOpa(OpaQueryInput input)
@@ -105,6 +125,31 @@ public <T> Set<T> parallelFilterFromOpa(
         return opaHttpClient.parallelFilterFromOpa(items, requestBuilder, opaPolicyUri, queryResultCodec);
     }
 
+    public List<OpaViewExpression> getRowFilterExpressionsFromOpa(OpaQueryContext context, CatalogSchemaTableName table)
+    {
+        OpaQueryInput queryInput = new OpaQueryInput(
+                context,
+                OpaQueryInputAction.builder()
+                        .operation("GetRowFilters")
+                        .resource(OpaQueryInputResource.builder().table(new TrinoTable(table)).build())
+                        .build());
+        return opaRowFiltersUri
+                .map(uri -> opaHttpClient.consumeOpaResponse(opaHttpClient.submitOpaRequest(queryInput, uri, rowFiltersQueryResultCodec)).result())
+                .orElse(ImmutableList.of());
+    }
+
+    public Optional<OpaViewExpression> getColumnMaskFromOpa(OpaQueryContext context, CatalogSchemaTableName table, String columnName, Type type)
+    {
+        OpaQueryInput queryInput = new OpaQueryInput(
+                context,
+                OpaQueryInputAction.builder()
+                        .operation("GetColumnMask")
+                        .resource(OpaQueryInputResource.builder().column(new TrinoColumn(table, columnName, type)).build())
+                        .build());
+        return opaColumnMaskingUri
+                .flatMap(uri -> opaHttpClient.consumeOpaResponse(opaHttpClient.submitOpaRequest(queryInput, uri, columnMaskQueryResultCodec)).result());
+    }
+
     public static OpaQueryInput buildQueryInputForSimpleResource(OpaQueryContext context, String operation, OpaQueryInputResource resource)
     {
         return new OpaQueryInput(context, OpaQueryInputAction.builder().operation(operation).resource(resource).build());

diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaColumnMaskQueryResult.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaColumnMaskQueryResult.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.opa.schema;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.validation.constraints.NotNull;
+
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public record OpaColumnMaskQueryResult(@JsonProperty("decision_id") String decisionId, @NotNull Optional<OpaViewExpression> result)
+{
+    public OpaColumnMaskQueryResult
+    {
+        requireNonNull(result, "result is null");
+    }
+}
diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaQueryInputResource.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaQueryInputResource.java
@@ -27,7 +27,8 @@ public record OpaQueryInputResource(
         TrinoFunction function,
         NamedEntity catalog,
         TrinoSchema schema,
-        TrinoTable table)
+        TrinoTable table,
+        TrinoColumn column)
 {
     public record NamedEntity(@NotNull String name)
     {
@@ -51,6 +52,7 @@ public static class Builder
         private TrinoSchema schema;
         private TrinoTable table;
         private TrinoFunction function;
+        private TrinoColumn column;
 
         private Builder() {}
 
@@ -102,6 +104,12 @@ public Builder function(String functionName)
             return this;
         }
 
+        public Builder column(TrinoColumn column)
+        {
+            this.column = column;
+            return this;
+        }
+
         public OpaQueryInputResource build()
         {
             return new OpaQueryInputResource(
@@ -111,7 +119,8 @@ public OpaQueryInputResource build()
                     this.function,
                     this.catalog,
                     this.schema,
-                    this.table);
+                    this.table,
+                    this.column);
         }
     }
 }
diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaRowFiltersQueryResult.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaRowFiltersQueryResult.java
@@ -0,0 +1,30 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.opa.schema;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+import jakarta.validation.constraints.NotNull;
+
+import java.util.List;
+
+import static java.util.Objects.requireNonNullElse;
+
+public record OpaRowFiltersQueryResult(@JsonProperty("decision_id") String decisionId, @NotNull List<OpaViewExpression> result)
+{
+    public OpaRowFiltersQueryResult
+    {
+        result = ImmutableList.copyOf(requireNonNullElse(result, ImmutableList.of()));
+    }
+}
diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaViewExpression.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaViewExpression.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.opa.schema;
+
+import io.trino.spi.security.ViewExpression;
+import jakarta.validation.constraints.NotNull;
+
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public record OpaViewExpression(@NotNull String expression, @NotNull Optional<String> identity)
+{
+    public OpaViewExpression
+    {
+        requireNonNull(expression, "expression is null");
+        requireNonNull(identity, "identity is null");
+    }
+
+    public ViewExpression toTrinoViewExpression(String catalogName, String schemaName)
+    {
+        ViewExpression.Builder builder = ViewExpression.builder()
+                .catalog(catalogName)
+                .schema(schemaName)
+                .expression(expression);
+        identity.ifPresent(builder::identity);
+        return builder.build();
+    }
+}