From e09f12ccb0c2d6ddac361f83edf7ecbe4c56638e Mon Sep 17 00:00:00 2001
From: Harshit <harshit.parikh@zomato.com>
Date: Mon, 21 Jul 2025 10:52:08 +0530
Subject: [PATCH] Add queryId to OPA requests

---
 .../sphinx/security/opa-access-control.md     |  4 ++
 .../io/trino/plugin/opa/OpaAccessControl.java |  4 +-
 .../plugin/opa/schema/OpaQueryContext.java    |  7 ++-
 .../plugin/opa/TestOpaAccessControl.java      | 47 +++++++++++++++++++
 4 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/docs/src/main/sphinx/security/opa-access-control.md b/docs/src/main/sphinx/security/opa-access-control.md
index 6a16bd9b6ba0..5df693dca5b3 100644
--- a/docs/src/main/sphinx/security/opa-access-control.md
+++ b/docs/src/main/sphinx/security/opa-access-control.md
@@ -130,6 +130,7 @@ The `context` object contains all other contextual information about the query:
   following two fields:
   - `user`: username
   - `groups`: list of groups this user belongs to
+- `queryId`: Query id
 - `softwareStack`: Information about the software stack issuing the request to
   OPA. The following information is included:
   - `trinoVersion`: Version of Trino used
@@ -158,6 +159,7 @@ Accessing a table results in a query similar to the following example:
       "user": "foo",
       "groups": ["some-group"]
     },
+    "queryId": "20250718_081710_03427_trino",
     "softwareStack": {
       "trinoVersion": "434"
     }
@@ -190,6 +192,7 @@ The `targetResource` is used in cases where a new resource, distinct from the on
       "user": "foo",
       "groups": ["some-group"]
     },
+    "queryId": "20250718_081710_03427_trino",
     "softwareStack": {
       "trinoVersion": "434"
     }
@@ -373,6 +376,7 @@ A batch column masking request is similar to the following example:
             "user": "foo",
             "groups": ["some-group"]
         },
+        "queryId": "20250718_081710_03427_trino",
         "softwareStack": {
             "trinoVersion": "434"
         }
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
index a0b92afee640..046a812c2c8c 100644
--- 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
@@ -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);
+        return new OpaQueryContext(TrinoIdentity.fromTrinoIdentity(trinoIdentity), pluginContext, Optional.empty());
     }
 
     OpaQueryContext buildQueryContext(SystemSecurityContext securityContext)
     {
-        return new OpaQueryContext(TrinoIdentity.fromTrinoIdentity(securityContext.getIdentity()), pluginContext);
+        return new OpaQueryContext(TrinoIdentity.fromTrinoIdentity(securityContext.getIdentity()), pluginContext, Optional.of(securityContext.getQueryId()));
     }
 }
diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaQueryContext.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaQueryContext.java
index 22c981135126..66fa7304af4a 100644
--- a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaQueryContext.java
+++ b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/schema/OpaQueryContext.java
@@ -13,13 +13,18 @@
  */
 package io.trino.plugin.opa.schema;
 
+import io.trino.spi.QueryId;
+
+import java.util.Optional;
+
 import static java.util.Objects.requireNonNull;
 
-public record OpaQueryContext(TrinoIdentity identity, OpaPluginContext softwareStack)
+public record OpaQueryContext(TrinoIdentity identity, OpaPluginContext softwareStack, Optional<QueryId> queryId)
 {
     public OpaQueryContext
     {
         requireNonNull(identity, "identity is null");
         requireNonNull(softwareStack, "softwareStack is null");
+        requireNonNull(queryId, "queryId is null");
     }
 }
diff --git a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControl.java b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControl.java
index 4740ebdf94e2..e2c795b4a92d 100644
--- a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControl.java
+++ b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControl.java
@@ -24,6 +24,7 @@
 import io.trino.plugin.opa.HttpClientUtils.InstrumentedHttpClient;
 import io.trino.plugin.opa.HttpClientUtils.MockResponse;
 import io.trino.plugin.opa.schema.OpaViewExpression;
+import io.trino.spi.QueryId;
 import io.trino.spi.connector.CatalogSchemaName;
 import io.trino.spi.connector.CatalogSchemaRoutineName;
 import io.trino.spi.connector.CatalogSchemaTableName;
@@ -37,6 +38,7 @@
 import io.trino.spi.type.VarcharType;
 import org.junit.jupiter.api.Test;
 
+import java.time.Instant;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -960,6 +962,51 @@ void testGetColumnMasksThrowsForIllegalResponse()
                 "Failed to deserialize");
     }
 
+
+    @Test
+    public void testQueryIdPropagation()
+    {
+        QueryId queryId = new QueryId("20250718_081710_03427_trino");
+
+        SystemSecurityContext customSecurityContext = new SystemSecurityContext(TEST_IDENTITY, queryId, Instant.now());
+        CatalogSchemaTableName tableName = new CatalogSchemaTableName("my_catalog", "my_schema", "my_table");
+
+        ThrowingMethodWrapper wrappedMethod = new ThrowingMethodWrapper(accessControl ->
+                accessControl.checkCanShowCreateTable(customSecurityContext, tableName));
+
+        String expectedActionRequest =
+                """
+                {
+                    "operation": "ShowCreateTable",
+                    "resource": {
+                        "table": {
+                            "catalogName": "%s",
+                            "schemaName": "%s",
+                            "tableName": "%s"
+                        }
+                    }
+                }
+                """.formatted(
+                        tableName.getCatalogName(),
+                        tableName.getSchemaTableName().getSchemaName(),
+                        tableName.getSchemaTableName().getTableName());
+
+        InstrumentedHttpClient mockClient = createMockHttpClient(OPA_SERVER_URI, request -> {
+            JsonNode contextNode = request.path("input").path("context");
+
+            assertThat(contextNode.path("queryId").asText()).isEqualTo(queryId.id());
+            assertThat(contextNode.path("identity").path("user").asText()).isEqualTo(TEST_IDENTITY.getUser());
+            assertThat(contextNode.path("softwareStack").path("trinoVersion").asText()).isEqualTo("trino-version");
+
+            return OK_RESPONSE;
+        });
+
+        OpaAccessControl authorizer = createOpaAuthorizer(simpleOpaConfig(), mockClient);
+
+        assertThat(wrappedMethod.isAccessAllowed(authorizer)).isTrue();
+        assertStringRequestsEqual(ImmutableSet.of(expectedActionRequest), mockClient.getRequests(), "/input/action");
+    }
+
     private void testGetColumnMasks(Map<ColumnSchema, String> columnResponseContent, Map<ColumnSchema, OpaViewExpression> expectedResult)
     {
         InstrumentedHttpClient httpClient = createMockHttpClient(