Cleanly reject session authorization queries for unsupported clients

Adds SESSION_AUTHORIZATION to ClientCapabilities, so that the client side support for this feature can be identified using the headers sent to Trino.
trinodb · Oct 18, 2023 · 60b9fc1 · 60b9fc1
1 parent 8b9f55e
commit 60b9fc1
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 1 deletion.
diff --git a/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java b/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java
@@ -23,5 +23,7 @@ public enum ClientCapabilities
     //   time(p) without time zone
     //   interval X(p1) to Y(p2)
     // When this capability is not set, the server returns datetime types with precision = 3
-    PARAMETRIC_DATETIME;
+    PARAMETRIC_DATETIME,
+    // Whether clients support the session authorization set/reset feature
+    SESSION_AUTHORIZATION;
 }
diff --git a/core/trino-main/src/main/java/io/trino/execution/ResetSessionAuthorizationTask.java b/core/trino-main/src/main/java/io/trino/execution/ResetSessionAuthorizationTask.java
@@ -16,6 +16,7 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.inject.Inject;
 import io.trino.Session;
+import io.trino.client.ClientCapabilities;
 import io.trino.execution.warnings.WarningCollector;
 import io.trino.spi.TrinoException;
 import io.trino.sql.tree.Expression;
@@ -26,6 +27,7 @@
 
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static io.trino.spi.StandardErrorCode.GENERIC_USER_ERROR;
+import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.util.Objects.requireNonNull;
 
 public class ResetSessionAuthorizationTask
@@ -53,6 +55,9 @@ public ListenableFuture<Void> execute(
             WarningCollector warningCollector)
     {
         Session session = stateMachine.getSession();
+        if (!session.getClientCapabilities().contains(ClientCapabilities.SESSION_AUTHORIZATION.toString())) {
+            throw new TrinoException(NOT_SUPPORTED, "RESET SESSION AUTHORIZATION not supported by client");
+        }
         session.getTransactionId().ifPresent(transactionId -> {
             if (!transactionManager.getTransactionInfo(transactionId).isAutoCommitContext()) {
                 throw new TrinoException(GENERIC_USER_ERROR, "Can't reset authorization user in the middle of a transaction");

diff --git a/core/trino-main/src/main/java/io/trino/execution/SetSessionAuthorizationTask.java b/core/trino-main/src/main/java/io/trino/execution/SetSessionAuthorizationTask.java
@@ -16,6 +16,7 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.inject.Inject;
 import io.trino.Session;
+import io.trino.client.ClientCapabilities;
 import io.trino.execution.warnings.WarningCollector;
 import io.trino.security.AccessControl;
 import io.trino.spi.TrinoException;
@@ -31,6 +32,7 @@
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static io.trino.spi.StandardErrorCode.GENERIC_USER_ERROR;
+import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.util.Objects.requireNonNull;
 
 public class SetSessionAuthorizationTask
@@ -60,6 +62,9 @@ public ListenableFuture<Void> execute(
             WarningCollector warningCollector)
     {
         Session session = stateMachine.getSession();
+        if (!session.getClientCapabilities().contains(ClientCapabilities.SESSION_AUTHORIZATION.toString())) {
+            throw new TrinoException(NOT_SUPPORTED, "SET SESSION AUTHORIZATION not supported by client");
+        }
         Identity originalIdentity = session.getOriginalIdentity();
         // Set authorization user in the middle of a transaction is disallowed by the SQL spec
         session.getTransactionId().ifPresent(transactionId -> {

diff --git a/testing/trino-tests/src/test/java/io/trino/tests/TestServer.java b/testing/trino-tests/src/test/java/io/trino/tests/TestServer.java
@@ -68,6 +68,7 @@
 import static io.trino.SystemSessionProperties.MAX_HASH_PARTITION_COUNT;
 import static io.trino.SystemSessionProperties.QUERY_MAX_MEMORY;
 import static io.trino.client.ClientCapabilities.PATH;
+import static io.trino.client.ClientCapabilities.SESSION_AUTHORIZATION;
 import static io.trino.client.ProtocolHeaders.TRINO_HEADERS;
 import static io.trino.spi.StandardErrorCode.INCOMPATIBLE_CLIENT;
 import static io.trino.testing.TestingSession.testSessionBuilder;
@@ -291,6 +292,34 @@ public void testSetPathSupportByClient()
         }
     }
 
+    @Test
+    public void testSetSessionSupportByClient()
+    {
+        try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of()).build())) {
+            assertThatThrownBy(() -> testingClient.execute("SET SESSION AUTHORIZATION userA"))
+                    .hasMessage("SET SESSION AUTHORIZATION not supported by client");
+        }
+
+        try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of(
+                SESSION_AUTHORIZATION.name())).build())) {
+            testingClient.execute("SET SESSION AUTHORIZATION userA");
+        }
+    }
+
+    @Test
+    public void testResetSessionSupportByClient()
+    {
+        try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of()).build())) {
+            assertThatThrownBy(() -> testingClient.execute("RESET SESSION AUTHORIZATION"))
+                    .hasMessage("RESET SESSION AUTHORIZATION not supported by client");
+        }
+
+        try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of(
+                SESSION_AUTHORIZATION.name())).build())) {
+            testingClient.execute("RESET SESSION AUTHORIZATION");
+        }
+    }
+
     private void checkVersionOnError(String query, @Language("RegExp") String proofOfOrigin)
     {
         QueryResults queryResults = postQuery(request -> request