Add query start to SystemSecurityContext

This can be used by SystemAccessCotrol implementors for providing visibility guarantees, i.e. to allow some sorts of caching.
trinodb · Sep 25, 2023 · 925f7e4 · 925f7e4
1 parent cb6c007
commit 925f7e4
Show file tree

Hide file tree

Showing 9 changed files with 114 additions and 89 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/Session.java b/core/trino-main/src/main/java/io/trino/Session.java
@@ -342,7 +342,7 @@ public Session beginTransactionId(TransactionId transactionId, TransactionManage
                 throw new TrinoException(NOT_FOUND, "Catalog for role does not exist: " + catalogName);
             }
             if (role.getType() == SelectedRole.Type.ROLE) {
-                accessControl.checkCanSetCatalogRole(new SecurityContext(transactionId, identity, queryId), role.getRole().orElseThrow(), catalogName);
+                accessControl.checkCanSetCatalogRole(new SecurityContext(transactionId, identity, queryId, start), role.getRole().orElseThrow(), catalogName);
             }
             connectorRoles.put(catalogName, role);
         }
@@ -559,7 +559,7 @@ private void validateCatalogProperties(
         for (Entry<String, String> property : catalogProperties.entrySet()) {
             // verify permissions
             if (transactionId.isPresent()) {
-                accessControl.checkCanSetCatalogSessionProperty(new SecurityContext(transactionId.get(), identity, queryId), catalogName, property.getKey());
+                accessControl.checkCanSetCatalogSessionProperty(new SecurityContext(transactionId.get(), identity, queryId, start), catalogName, property.getKey());
             }
 
             // validate catalog session property value
@@ -591,7 +591,7 @@ public static SessionBuilder builder(Session session)
 
     public SecurityContext toSecurityContext()
     {
-        return new SecurityContext(getRequiredTransactionId(), getIdentity(), queryId);
+        return new SecurityContext(getRequiredTransactionId(), getIdentity(), queryId, start);
     }
 
     public static class SessionBuilder

diff --git a/core/trino-main/src/main/java/io/trino/security/SecurityContext.java b/core/trino-main/src/main/java/io/trino/security/SecurityContext.java
@@ -19,6 +19,7 @@
 import io.trino.spi.security.SystemSecurityContext;
 import io.trino.transaction.TransactionId;
 
+import java.time.Instant;
 import java.util.Objects;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
@@ -29,18 +30,20 @@ public class SecurityContext
     public static SecurityContext of(Session session)
     {
         requireNonNull(session, "session is null");
-        return new SecurityContext(session.getRequiredTransactionId(), session.getIdentity(), session.getQueryId());
+        return new SecurityContext(session.getRequiredTransactionId(), session.getIdentity(), session.getQueryId(), session.getStart());
     }
 
     private final TransactionId transactionId;
     private final Identity identity;
     private final QueryId queryId;
+    private final Instant queryStart;
 
-    public SecurityContext(TransactionId transactionId, Identity identity, QueryId queryId)
+    public SecurityContext(TransactionId transactionId, Identity identity, QueryId queryId, Instant queryStart)
     {
         this.transactionId = requireNonNull(transactionId, "transactionId is null");
         this.identity = requireNonNull(identity, "identity is null");
         this.queryId = requireNonNull(queryId, "queryId is null");
+        this.queryStart = requireNonNull(queryStart, "queryStart is null");
     }
 
     public TransactionId getTransactionId()
@@ -60,7 +63,7 @@ public QueryId getQueryId()
 
     public SystemSecurityContext toSystemSecurityContext()
     {
-        return new SystemSecurityContext(identity, queryId);
+        return new SystemSecurityContext(identity, queryId, queryStart);
     }
 
     @Override

diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/Analysis.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/Analysis.java
@@ -87,6 +87,7 @@
 import io.trino.transaction.TransactionId;
 import jakarta.annotation.Nullable;
 
+import java.time.Instant;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -1841,9 +1842,9 @@ public AccessControl getAccessControl()
             return accessControl;
         }
 
-        public SecurityContext getSecurityContext(TransactionId transactionId, QueryId queryId)
+        public SecurityContext getSecurityContext(TransactionId transactionId, QueryId queryId, Instant queryStart)
         {
-            return new SecurityContext(transactionId, identity, queryId);
+            return new SecurityContext(transactionId, identity, queryId, queryStart);
         }
 
         @Override

diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/Analyzer.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/Analyzer.java
@@ -101,7 +101,7 @@ public Analysis analyze(Statement statement, QueryType queryType)
             analysis.getTableColumnReferences().forEach((accessControlInfo, tableColumnReferences) ->
                     tableColumnReferences.forEach((tableName, columns) ->
                             accessControlInfo.getAccessControl().checkCanSelectFromColumns(
-                                    accessControlInfo.getSecurityContext(session.getRequiredTransactionId(), session.getQueryId()),
+                                    accessControlInfo.getSecurityContext(session.getRequiredTransactionId(), session.getQueryId(), session.getStart()),
                                     tableName,
                                     columns)));
         }

diff --git a/core/trino-main/src/test/java/io/trino/security/TestAccessControlManager.java b/core/trino-main/src/test/java/io/trino/security/TestAccessControlManager.java
@@ -51,6 +51,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.security.Principal;
+import java.time.Instant;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -77,6 +78,7 @@ public class TestAccessControlManager
     private static final Principal PRINCIPAL = new BasicPrincipal("principal");
     private static final String USER_NAME = "user_name";
     private static final QueryId queryId = new QueryId("query_id");
+    private static final Instant queryStart = Instant.now();
 
     @Test
     public void testInitializing()
@@ -109,7 +111,7 @@ public void testReadOnlySystemAccessControl()
 
         transaction(transactionManager, accessControlManager)
                 .execute(transactionId -> {
-                    SecurityContext context = new SecurityContext(transactionId, identity, queryId);
+                    SecurityContext context = new SecurityContext(transactionId, identity, queryId, queryStart);
                     accessControlManager.checkCanSetCatalogSessionProperty(context, TEST_CATALOG_NAME, "property");
                     accessControlManager.checkCanShowSchemas(context, TEST_CATALOG_NAME);
                     accessControlManager.checkCanShowTables(context, new CatalogSchemaName(TEST_CATALOG_NAME, "schema"));
@@ -127,7 +129,7 @@ public void testReadOnlySystemAccessControl()
 
         assertThatThrownBy(() -> transaction(transactionManager, accessControlManager)
                 .execute(transactionId -> {
-                    accessControlManager.checkCanInsertIntoTable(new SecurityContext(transactionId, identity, queryId), tableName);
+                    accessControlManager.checkCanInsertIntoTable(new SecurityContext(transactionId, identity, queryId, queryStart), tableName);
                 }))
                 .isInstanceOf(AccessDeniedException.class)
                 .hasMessage("Access Denied: Cannot insert into table test-catalog.schema.table");
@@ -264,7 +266,7 @@ public void checkCanShowCreateTable(ConnectorSecurityContext context, SchemaTabl
     private static SecurityContext context(TransactionId transactionId)
     {
         Identity identity = Identity.forUser(USER_NAME).withPrincipal(PRINCIPAL).build();
-        return new SecurityContext(transactionId, identity, queryId);
+        return new SecurityContext(transactionId, identity, queryId, queryStart);
     }
 
     @Test