[OPA] Add additional context file to Trino OPA plugin

docs/src/main/sphinx/security/opa-access-control.md

@@ -65,6 +65,10 @@ The following table lists the configuration properties for the OPA access contro
   - Optional HTTP client configurations for the connection from Trino to OPA,
     for example `opa.http-client.http-proxy` for configuring the HTTP proxy.
     Find more details in [](/admin/properties-http-client).
+* - `opa.context-file`
+  - Optional properties file, containing user defined properties
+    (e.g. tenant namespace, tier or cluster) to be included in
+    the OPA query context.
 :::
 
 ### Logging

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

@@ -802,11 +802,11 @@ private static Map<String, Optional<Object>> convertProperties(Map<String, Objec
 
     OpaQueryContext buildQueryContext(Identity trinoIdentity)
     {
-        return new OpaQueryContext(TrinoIdentity.fromTrinoIdentity(trinoIdentity), pluginContext, Optional.empty());
+        return new OpaQueryContext(TrinoIdentity.fromTrinoIdentity(trinoIdentity), pluginContext, opaHighLevelClient.getAdditionalContext(), Optional.empty());
     }
 
     OpaQueryContext buildQueryContext(SystemSecurityContext securityContext)
     {
-        return new OpaQueryContext(TrinoIdentity.fromTrinoIdentity(securityContext.getIdentity()), pluginContext, Optional.of(securityContext.getQueryId()));
+        return new OpaQueryContext(TrinoIdentity.fromTrinoIdentity(securityContext.getIdentity()), pluginContext, opaHighLevelClient.getAdditionalContext(), Optional.of(securityContext.getQueryId()));
     }
 }

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

@@ -15,9 +15,11 @@
 
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
+import io.airlift.configuration.validation.FileExists;
 import jakarta.validation.constraints.NotNull;
 
 import java.net.URI;
+import java.nio.file.Path;
 import java.util.Optional;
 
 public class OpaConfig
@@ -31,6 +33,7 @@ public class OpaConfig
     private Optional<URI> opaRowFiltersUri = Optional.empty();
     private Optional<URI> opaColumnMaskingUri = Optional.empty();
     private Optional<URI> opaBatchColumnMaskingUri = Optional.empty();
+    private Optional<Path> additionalContextFile = Optional.empty();
 
     @NotNull
     public URI getOpaUri()
@@ -140,4 +143,17 @@ public OpaConfig setOpaBatchColumnMaskingUri(URI opaBatchColumnMaskingUri)
         this.opaBatchColumnMaskingUri = Optional.ofNullable(opaBatchColumnMaskingUri);
         return this;
     }
+
+    public Optional<@FileExists Path> getAdditionalContextFile()
+    {
+        return additionalContextFile;
+    }
+
+    @Config("opa.context-file")
+    @ConfigDescription("Optional properties file, containing user defined properties (e.g. tenant namespace, tier or cluster) to be included in the OPA query context")
+    public OpaConfig setAdditionalContextFile(Path additionalContextFile)
+    {
+        this.additionalContextFile = Optional.ofNullable(additionalContextFile);
+        return this;
+    }
 }

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

@@ -32,16 +32,22 @@
 import io.trino.spi.connector.ColumnSchema;
 import io.trino.spi.security.AccessDeniedException;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
 import java.net.URI;
+import java.nio.file.Path;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Properties;
 import java.util.Set;
 import java.util.function.Function;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static java.nio.file.Files.newInputStream;
 import static java.util.Objects.requireNonNull;
 
 public class OpaHighLevelClient
@@ -55,6 +61,7 @@ public class OpaHighLevelClient
     private final Optional<URI> opaRowFiltersUri;
     private final Optional<URI> opaColumnMaskingUri;
     private final Optional<URI> opaBatchColumnMaskingUri;
+    private final ImmutableMap<String, String> opaAdditionalContext;
 
     @Inject
     public OpaHighLevelClient(
@@ -74,6 +81,7 @@ public OpaHighLevelClient(
         this.opaRowFiltersUri = config.getOpaRowFiltersUri();
         this.opaColumnMaskingUri = config.getOpaColumnMaskingUri();
         this.opaBatchColumnMaskingUri = config.getOpaBatchColumnMaskingUri();
+        this.opaAdditionalContext = ImmutableMap.copyOf(loadAdditionalContextFromFile(config.getAdditionalContextFile()));
     }
 
     public boolean queryOpa(OpaQueryInput input)
@@ -155,6 +163,11 @@ public Map<ColumnSchema, OpaViewExpression> getColumnMasksFromOpa(OpaQueryContex
                 .orElse(ImmutableMap.of());
     }
 
+    public Map<String, String> getAdditionalContext()
+    {
+        return opaAdditionalContext;
+    }
+
     public static OpaQueryInput buildQueryInputForSimpleResource(OpaQueryContext context, String operation, OpaQueryInputResource resource)
     {
         return new OpaQueryInput(context, OpaQueryInputAction.builder().operation(operation).resource(resource).build());
@@ -190,4 +203,23 @@ private static OpaQueryInput buildQueryInputForSimpleAction(OpaQueryContext cont
     {
         return new OpaQueryInput(context, OpaQueryInputAction.builder().operation(operation).build());
     }
+
+    private static Map<String, String> loadAdditionalContextFromFile(Optional<Path> additionalContextFile)
+    {
+        if (additionalContextFile.isEmpty()) {
+            return ImmutableMap.of();
+        }
+
+        try (InputStream inputStream = newInputStream(additionalContextFile.get())) {
+            Properties properties = new Properties();
+            properties.load(inputStream);
+            return properties.entrySet().stream()
+                    .collect(toImmutableMap(
+                            entry -> String.valueOf(entry.getKey()),
+                            entry -> String.valueOf(entry.getValue())));
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
 }

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

@@ -13,18 +13,21 @@
  */
 package io.trino.plugin.opa.schema;
 
+import com.google.common.collect.ImmutableMap;
 import io.trino.spi.QueryId;
 
+import java.util.Map;
 import java.util.Optional;
 
 import static java.util.Objects.requireNonNull;
 
-public record OpaQueryContext(TrinoIdentity identity, OpaPluginContext softwareStack, Optional<QueryId> queryId)
+public record OpaQueryContext(TrinoIdentity identity, OpaPluginContext softwareStack, Map<String, String> properties, Optional<QueryId> queryId)
 {
     public OpaQueryContext
     {
         requireNonNull(identity, "identity is null");
         requireNonNull(softwareStack, "softwareStack is null");
+        properties = ImmutableMap.copyOf(properties);
         requireNonNull(queryId, "queryId is null");
     }
 }

plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestConstants.java

@@ -23,6 +23,8 @@
 import io.trino.spi.security.SystemSecurityContext;
 
 import java.net.URI;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.time.Instant;
 
 public final class TestConstants
@@ -58,6 +60,7 @@ private TestConstants() {}
     public static final URI OPA_SERVER_BATCH_URI = URI.create("http://my-batch-uri/");
     public static final URI OPA_ROW_FILTERING_URI = URI.create("http://my-row-filtering-uri/");
     public static final URI OPA_COLUMN_MASKING_URI = URI.create("http://my-column-masking-uri/");
+    public static final Path OPA_ADDITIONAL_CONTEXT_FILE = Paths.get("src/test/resources/additional-context.properties");
     public static final Identity TEST_IDENTITY = Identity.forUser("source-user").withGroups(ImmutableSet.of("some-group")).build();
     public static final QueryId TEST_QUERY_ID = QueryId.valueOf("abcde");
     public static final SystemSecurityContext TEST_SECURITY_CONTEXT = new SystemSecurityContext(TEST_IDENTITY, new QueryIdGenerator().createNextQueryId(), Instant.now());

plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestHelpers.java

@@ -92,6 +92,7 @@ public static Map<String, String> opaConfigToDict(OpaConfig config)
         config.getOpaRowFiltersUri().ifPresent(rowFiltersUri -> configBuilder.put("opa.policy.row-filters-uri", rowFiltersUri.toString()));
         config.getOpaColumnMaskingUri().ifPresent(columnMaskingUri -> configBuilder.put("opa.policy.column-masking-uri", columnMaskingUri.toString()));
         config.getOpaBatchColumnMaskingUri().ifPresent(batchColumnMaskingUri -> configBuilder.put("opa.policy.batch-column-masking-uri", batchColumnMaskingUri.toString()));
+        config.getAdditionalContextFile().ifPresent(additionalContextFile -> configBuilder.put("opa.context-file", additionalContextFile.toString()));
         return configBuilder.buildOrThrow();
     }
 

plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControl.java

@@ -36,6 +36,7 @@
 import io.trino.spi.security.TrinoPrincipal;
 import io.trino.spi.security.ViewExpression;
 import io.trino.spi.type.VarcharType;
+import org.intellij.lang.annotations.Language;
 import org.junit.jupiter.api.Test;
 
 import java.time.Instant;
@@ -57,6 +58,7 @@
 import static io.trino.plugin.opa.TestConstants.MALFORMED_RESPONSE;
 import static io.trino.plugin.opa.TestConstants.NO_ACCESS_RESPONSE;
 import static io.trino.plugin.opa.TestConstants.OK_RESPONSE;
+import static io.trino.plugin.opa.TestConstants.OPA_ADDITIONAL_CONTEXT_FILE;
 import static io.trino.plugin.opa.TestConstants.OPA_COLUMN_MASKING_URI;
 import static io.trino.plugin.opa.TestConstants.OPA_ROW_FILTERING_URI;
 import static io.trino.plugin.opa.TestConstants.OPA_SERVER_URI;
@@ -682,13 +684,52 @@ private void testRequestContextContentsForGivenTrinoVersion(Optional<SystemAcces
                         },
                         "softwareStack": {
                             "trinoVersion": "%s"
+                        },
+                        "properties" : {
                         }
                     }
                 }\
                 """.formatted(expectedTrinoVersion);
         assertStringRequestsEqual(ImmutableSet.of(expectedRequest), mockClient.getRequests(), "/input");
     }
 
+    @Test
+    void testRequestContextContentsForAdditionalContext()
+    {
+        InstrumentedHttpClient mockClient = createMockHttpClient(OPA_SERVER_URI, request -> OK_RESPONSE);
+        OpaAccessControl authorizer = (OpaAccessControl) OpaAccessControlFactory.create(
+                ImmutableMap.of("opa.policy.uri", OPA_SERVER_URI.toString(), "opa.context-file", OPA_ADDITIONAL_CONTEXT_FILE.toString()),
+                Optional.of(mockClient),
+                Optional.empty());
+        Identity sampleIdentityWithGroups = Identity.forUser("test_user").withGroups(ImmutableSet.of("some_group")).build();
+
+        authorizer.checkCanExecuteQuery(sampleIdentityWithGroups, TEST_QUERY_ID);
+
+        @Language("JSON")
+        String expectedRequest =
+                """
+                {
+                    "action": {
+                        "operation": "ExecuteQuery"
+                    },
+                    "context": {
+                        "identity": {
+                            "user": "test_user",
+                            "groups": ["some_group"]
+                        },
+                        "softwareStack": {
+                            "trinoVersion": "UNKNOWN"
+                        },
+                        "properties" : {
+                            "namespace" : "some-namespace",
+                            "cluster" : "some-cluster"
+                        }
+                    }
+                }
+                """;
+        assertStringRequestsEqual(ImmutableSet.of(expectedRequest), mockClient.getRequests(), "/input");
+    }
+
     @Test
     void testGetRowFiltersThrowsForIllegalResponse()
     {