Rebase and implement new SPI methods

trinodb · Oct 27, 2023 · 3aec31d · 3aec31d
1 parent c3e8ded
commit 3aec31d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
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
@@ -53,12 +53,14 @@
 import static io.trino.plugin.opa.OpaHighLevelClient.buildQueryInputForSimpleResource;
 import static io.trino.plugin.opa.schema.PropertiesMapper.convertProperties;
 import static io.trino.spi.security.AccessDeniedException.denyCreateCatalog;
+import static io.trino.spi.security.AccessDeniedException.denyCreateFunction;
 import static io.trino.spi.security.AccessDeniedException.denyCreateRole;
 import static io.trino.spi.security.AccessDeniedException.denyCreateSchema;
 import static io.trino.spi.security.AccessDeniedException.denyCreateViewWithSelect;
 import static io.trino.spi.security.AccessDeniedException.denyDenySchemaPrivilege;
 import static io.trino.spi.security.AccessDeniedException.denyDenyTablePrivilege;
 import static io.trino.spi.security.AccessDeniedException.denyDropCatalog;
+import static io.trino.spi.security.AccessDeniedException.denyDropFunction;
 import static io.trino.spi.security.AccessDeniedException.denyDropRole;
 import static io.trino.spi.security.AccessDeniedException.denyDropSchema;
 import static io.trino.spi.security.AccessDeniedException.denyExecuteProcedure;
@@ -840,6 +842,26 @@ public void checkCanExecuteTableProcedure(SystemSecurityContext systemSecurityCo
                 OpaQueryInputResource.builder().table(new TrinoTable(table)).function(procedure).build());
     }
 
+    @Override
+    public void checkCanCreateFunction(SystemSecurityContext systemSecurityContext, CatalogSchemaRoutineName functionName)
+    {
+        opaHighLevelClient.queryAndEnforce(
+                OpaQueryContext.fromSystemSecurityContext(systemSecurityContext),
+                "CreateFunction",
+                () -> denyCreateFunction(functionName.toString()),
+                OpaQueryInputResource.builder().function(TrinoFunction.fromTrinoFunction(functionName)).build());
+    }
+
+    @Override
+    public void checkCanDropFunction(SystemSecurityContext systemSecurityContext, CatalogSchemaRoutineName functionName)
+    {
+        opaHighLevelClient.queryAndEnforce(
+                OpaQueryContext.fromSystemSecurityContext(systemSecurityContext),
+                "DropFunction",
+                () -> denyDropFunction(functionName.toString()),
+                OpaQueryInputResource.builder().function(TrinoFunction.fromTrinoFunction(functionName)).build());
+    }
+
     private void checkTableOperation(SystemSecurityContext context, String actionName, CatalogSchemaTableName table, Consumer<String> deny)
     {
         opaHighLevelClient.queryAndEnforce(

diff --git a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/OpaAccessControlUnitTest.java b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/OpaAccessControlUnitTest.java
@@ -1324,10 +1324,14 @@ private static Stream<Arguments> functionResourceTestCases()
     {
         Stream<TestHelpers.MethodWrapper<CatalogSchemaRoutineName>> methods = Stream.of(
                 new TestHelpers.ThrowingMethodWrapper<>(OpaAccessControl::checkCanExecuteProcedure),
+                new TestHelpers.ThrowingMethodWrapper<>(OpaAccessControl::checkCanCreateFunction),
+                new TestHelpers.ThrowingMethodWrapper<>(OpaAccessControl::checkCanDropFunction),
                 new TestHelpers.ReturningMethodWrapper<>(OpaAccessControl::canExecuteFunction),
                 new TestHelpers.ReturningMethodWrapper<>(OpaAccessControl::canCreateViewWithExecuteFunction));
         Stream<String> actions = Stream.of(
                 "ExecuteProcedure",
+                "CreateFunction",
+                "DropFunction",
                 "ExecuteFunction",
                 "CreateViewWithExecuteFunction");
         return Streams.zip(actions, methods, (action, method) -> Arguments.of(Named.of(action, action), method));