Record routines in query event

presto-main/src/main/java/io/prestosql/dispatcher/FailedDispatchQuery.java

@@ -227,6 +227,7 @@ private static QueryInfo immediateFailureQueryInfo(
                 ImmutableSet.of(),
                 Optional.empty(),
                 ImmutableList.of(),
+                ImmutableList.of(),
                 true,
                 resourceGroupId);
 

presto-main/src/main/java/io/prestosql/event/QueryMonitor.java

@@ -126,6 +126,7 @@ public void queryCreatedEvent(BasicQueryInfo queryInfo)
                                 queryInfo.getPreparedQuery(),
                                 QUEUED.toString(),
                                 ImmutableList.of(),
+                                ImmutableList.of(),
                                 queryInfo.getSelf(),
                                 Optional.empty(),
                                 Optional.empty())));
@@ -141,6 +142,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                         queryInfo.getPreparedQuery(),
                         queryInfo.getState().toString(),
                         ImmutableList.of(),
+                        ImmutableList.of(),
                         queryInfo.getSelf(),
                         Optional.empty(),
                         Optional.empty()),
@@ -210,6 +212,7 @@ private QueryMetadata createQueryMetadata(QueryInfo queryInfo)
                 queryInfo.getPreparedQuery(),
                 queryInfo.getState().toString(),
                 queryInfo.getReferencedTables(),
+                queryInfo.getRoutines(),
                 queryInfo.getSelf(),
                 createTextQueryPlan(queryInfo),
                 queryInfo.getOutputStage().flatMap(stage -> stageInfoCodec.toJsonWithLengthLimit(stage, maxJsonLimit)));

presto-main/src/main/java/io/prestosql/execution/CallTask.java

@@ -13,6 +13,7 @@
  */
 package io.prestosql.execution;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
 import io.prestosql.Session;
 import io.prestosql.connector.CatalogName;
@@ -22,6 +23,7 @@
 import io.prestosql.spi.PrestoException;
 import io.prestosql.spi.block.BlockBuilder;
 import io.prestosql.spi.connector.ConnectorSession;
+import io.prestosql.spi.eventlistener.RoutineInfo;
 import io.prestosql.spi.procedure.Procedure;
 import io.prestosql.spi.procedure.Procedure.Argument;
 import io.prestosql.spi.type.Type;
@@ -173,6 +175,7 @@ else if (i < procedure.getArguments().size()) {
         }
 
         accessControl.checkCanExecuteProcedure(session.toSecurityContext(), procedureName);
+        stateMachine.setRoutines(ImmutableList.of(new RoutineInfo(procedureName.getObjectName(), session.getUser())));
 
         try {
             procedure.getMethodHandle().invokeWithArguments(arguments);

presto-main/src/main/java/io/prestosql/execution/QueryInfo.java

@@ -23,6 +23,7 @@
 import io.prestosql.spi.ErrorType;
 import io.prestosql.spi.PrestoWarning;
 import io.prestosql.spi.QueryId;
+import io.prestosql.spi.eventlistener.RoutineInfo;
 import io.prestosql.spi.eventlistener.TableInfo;
 import io.prestosql.spi.memory.MemoryPoolId;
 import io.prestosql.spi.resourcegroups.ResourceGroupId;
@@ -69,6 +70,7 @@ public class QueryInfo
     private final String updateType;
     private final Optional<StageInfo> outputStage;
     private final List<TableInfo> referencedTables;
+    private final List<RoutineInfo> routines;
     private final ExecutionFailureInfo failureInfo;
     private final ErrorType errorType;
     private final ErrorCode errorCode;
@@ -108,6 +110,7 @@ public QueryInfo(
             @JsonProperty("inputs") Set<Input> inputs,
             @JsonProperty("output") Optional<Output> output,
             @JsonProperty("referencedTables") List<TableInfo> referencedTables,
+            @JsonProperty("routines") List<RoutineInfo> routines,
             @JsonProperty("completeInfo") boolean completeInfo,
             @JsonProperty("resourceGroupId") Optional<ResourceGroupId> resourceGroupId)
     {
@@ -131,6 +134,7 @@ public QueryInfo(
         requireNonNull(inputs, "inputs is null");
         requireNonNull(output, "output is null");
         requireNonNull(referencedTables, "referencedTables is null");
+        requireNonNull(routines, "routines is null");
         requireNonNull(resourceGroupId, "resourceGroupId is null");
         requireNonNull(warnings, "warnings is null");
 
@@ -163,6 +167,7 @@ public QueryInfo(
         this.inputs = ImmutableSet.copyOf(inputs);
         this.output = output;
         this.referencedTables = ImmutableList.copyOf(referencedTables);
+        this.routines = ImmutableList.copyOf(routines);
         this.completeInfo = completeInfo;
         this.resourceGroupId = resourceGroupId;
     }
@@ -351,6 +356,12 @@ public List<TableInfo> getReferencedTables()
         return referencedTables;
     }
 
+    @JsonProperty
+    public List<RoutineInfo> getRoutines()
+    {
+        return routines;
+    }
+
     @JsonProperty
     public Optional<ResourceGroupId> getResourceGroupId()
     {

presto-main/src/main/java/io/prestosql/execution/QueryStateMachine.java

@@ -38,6 +38,7 @@
 import io.prestosql.spi.ErrorCode;
 import io.prestosql.spi.PrestoException;
 import io.prestosql.spi.QueryId;
+import io.prestosql.spi.eventlistener.RoutineInfo;
 import io.prestosql.spi.eventlistener.StageGcStatistics;
 import io.prestosql.spi.eventlistener.TableInfo;
 import io.prestosql.spi.resourcegroups.ResourceGroupId;
@@ -150,6 +151,7 @@ public class QueryStateMachine
     private final AtomicReference<Set<Input>> inputs = new AtomicReference<>(ImmutableSet.of());
     private final AtomicReference<Optional<Output>> output = new AtomicReference<>(Optional.empty());
     private final AtomicReference<List<TableInfo>> referencedTables = new AtomicReference<>(ImmutableList.of());
+    private final AtomicReference<List<RoutineInfo>> routines = new AtomicReference<>(ImmutableList.of());
     private final StateMachine<Optional<QueryInfo>> finalQueryInfo;
 
     private final WarningCollector warningCollector;
@@ -424,6 +426,7 @@ QueryInfo getQueryInfo(Optional<StageInfo> rootStage)
                 inputs.get(),
                 output.get(),
                 referencedTables.get(),
+                routines.get(),
                 completeInfo,
                 Optional.of(resourceGroup));
     }
@@ -636,6 +639,12 @@ public void setReferencedTables(List<TableInfo> tables)
         referencedTables.set(ImmutableList.copyOf(tables));
     }
 
+    public void setRoutines(List<RoutineInfo> routines)
+    {
+        requireNonNull(routines, "routines is null");
+        this.routines.set(ImmutableList.copyOf(routines));
+    }
+
     public Map<String, String> getSetSessionProperties()
     {
         return setSessionProperties;
@@ -1041,6 +1050,7 @@ public void pruneQueryInfo()
                 queryInfo.getInputs(),
                 queryInfo.getOutput(),
                 queryInfo.getReferencedTables(),
+                queryInfo.getRoutines(),
                 queryInfo.isCompleteInfo(),
                 queryInfo.getResourceGroupId());
         finalQueryInfo.compareAndSet(finalInfo, Optional.of(prunedQueryInfo));

presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java

@@ -221,6 +221,7 @@ private Analysis analyze(
 
         stateMachine.setUpdateType(analysis.getUpdateType());
         stateMachine.setReferencedTables(analysis.getReferencedTables());
+        stateMachine.setRoutines(analysis.getRoutines());
 
         stateMachine.endAnalysis();
 

presto-main/src/main/java/io/prestosql/sql/analyzer/Analysis.java

@@ -30,6 +30,7 @@
 import io.prestosql.spi.connector.ColumnHandle;
 import io.prestosql.spi.connector.ConnectorTableMetadata;
 import io.prestosql.spi.eventlistener.ColumnInfo;
+import io.prestosql.spi.eventlistener.RoutineInfo;
 import io.prestosql.spi.eventlistener.TableInfo;
 import io.prestosql.spi.security.Identity;
 import io.prestosql.spi.security.ViewExpression;
@@ -134,7 +135,7 @@ public class Analysis
     private final Map<NodeRef<Expression>, Type> coercions = new LinkedHashMap<>();
     private final Set<NodeRef<Expression>> typeOnlyCoercions = new LinkedHashSet<>();
     private final Map<NodeRef<Relation>, List<Type>> relationCoercions = new LinkedHashMap<>();
-    private final Map<NodeRef<FunctionCall>, ResolvedFunction> resolvedFunctions = new LinkedHashMap<>();
+    private final Map<NodeRef<FunctionCall>, RoutineEntry> resolvedFunctions = new LinkedHashMap<>();
     private final Map<NodeRef<Identifier>, LambdaArgumentDeclaration> lambdaArgumentReferences = new LinkedHashMap<>();
 
     private final Map<Field, ColumnHandle> columns = new LinkedHashMap<>();
@@ -514,12 +515,12 @@ public void registerTable(
 
     public ResolvedFunction getResolvedFunction(FunctionCall function)
     {
-        return resolvedFunctions.get(NodeRef.of(function));
+        return resolvedFunctions.get(NodeRef.of(function)).getFunction();
     }
 
-    public void addResolvedFunction(Map<NodeRef<FunctionCall>, ResolvedFunction> infos)
+    public void addResolvedFunction(FunctionCall node, ResolvedFunction function, String authorization)
     {
-        resolvedFunctions.putAll(infos);
+        resolvedFunctions.put(NodeRef.of(node), new RoutineEntry(function, authorization));
     }
 
     public Set<NodeRef<Expression>> getColumnReferences()
@@ -794,6 +795,13 @@ public List<TableInfo> getReferencedTables()
                 .collect(toImmutableList());
     }
 
+    public List<RoutineInfo> getRoutines()
+    {
+        return resolvedFunctions.entrySet().stream()
+                .map(entry -> new RoutineInfo(entry.getValue().function.getSignature().getName(), entry.getValue().getAuthorization()))
+                .collect(toImmutableList());
+    }
+
     @Immutable
     public static final class SelectExpression
     {
@@ -1136,4 +1144,26 @@ public String getAuthorization()
             return authorization;
         }
     }
+
+    private static class RoutineEntry
+    {
+        private final ResolvedFunction function;
+        private final String authorization;
+
+        public RoutineEntry(ResolvedFunction function, String authorization)
+        {
+            this.function = requireNonNull(function, "function is null");
+            this.authorization = requireNonNull(authorization, "authorization is null");
+        }
+
+        public ResolvedFunction getFunction()
+        {
+            return function;
+        }
+
+        public String getAuthorization()
+        {
+            return authorization;
+        }
+    }
 }

presto-main/src/main/java/io/prestosql/sql/analyzer/ExpressionAnalyzer.java

@@ -1598,7 +1598,10 @@ public static ExpressionAnalysis analyzeExpression(
 
         analysis.addTypes(expressionTypes);
         analysis.addCoercions(expressionCoercions, typeOnlyCoercions);
-        analysis.addResolvedFunction(resolvedFunctions);
+
+        resolvedFunctions.entrySet()
+                .forEach(entry -> analysis.addResolvedFunction(entry.getKey().getNode(), entry.getValue(), session.getUser()));
+
         analysis.addColumnReferences(analyzer.getColumnReferences());
         analysis.addLambdaArgumentReferences(analyzer.getLambdaArgumentReferences());
         analysis.addTableColumnReferences(accessControl, session.getIdentity(), analyzer.getTableColumnReferences());

presto-main/src/test/java/io/prestosql/execution/MockManagedQueryExecution.java

@@ -237,6 +237,7 @@ public QueryInfo getFullQueryInfo()
                 ImmutableSet.of(),
                 Optional.empty(),
                 ImmutableList.of(),
+                ImmutableList.of(),
                 state.isDone(),
                 Optional.empty());
     }

presto-main/src/test/java/io/prestosql/server/TestBasicQueryInfo.java

@@ -128,6 +128,7 @@ public void testConstructor()
                         ImmutableSet.of(),
                         Optional.empty(),
                         ImmutableList.of(),
+                        ImmutableList.of(),
                         false,
                         Optional.empty()));
 

presto-main/src/test/java/io/prestosql/server/TestQueryStateInfo.java

@@ -172,6 +172,7 @@ private QueryInfo createQueryInfo(String queryId, QueryState state, String query
                 ImmutableSet.of(),
                 Optional.empty(),
                 ImmutableList.of(),
+                ImmutableList.of(),
                 false,
                 Optional.empty());
     }

presto-spi/src/main/java/io/prestosql/spi/eventlistener/QueryMetadata.java

@@ -33,6 +33,7 @@ public class QueryMetadata
     private final URI uri;
 
     private final List<TableInfo> tables;
+    private final List<RoutineInfo> routines;
 
     private final Optional<String> plan;
 
@@ -45,6 +46,7 @@ public QueryMetadata(
             Optional<String> preparedQuery,
             String queryState,
             List<TableInfo> tables,
+            List<RoutineInfo> routines,
             URI uri,
             Optional<String> plan,
             Optional<String> payload)
@@ -55,6 +57,7 @@ public QueryMetadata(
         this.preparedQuery = requireNonNull(preparedQuery, "preparedQuery is null");
         this.queryState = requireNonNull(queryState, "queryState is null");
         this.tables = requireNonNull(tables, "tables is null");
+        this.routines = requireNonNull(routines, "routines is null");
         this.uri = requireNonNull(uri, "uri is null");
         this.plan = requireNonNull(plan, "plan is null");
         this.payload = requireNonNull(payload, "payload is null");
@@ -96,6 +99,12 @@ public List<TableInfo> getTables()
         return tables;
     }
 
+    @JsonProperty
+    public List<RoutineInfo> getRoutines()
+    {
+        return routines;
+    }
+
     @JsonProperty
     public URI getUri()
     {

presto-spi/src/main/java/io/prestosql/spi/eventlistener/RoutineInfo.java

@@ -13,6 +13,34 @@
  */
 package io.prestosql.spi.eventlistener;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import static java.util.Objects.requireNonNull;
+
 public class RoutineInfo
 {
+    private final String routine;
+    private final String authorization;
+
+    @JsonCreator
+    public RoutineInfo(
+            @JsonProperty("routine") String routine,
+            @JsonProperty("authorization") String authorization)
+    {
+        this.routine = requireNonNull(routine, "routine is null");
+        this.authorization = requireNonNull(authorization, "authorization is null");
+    }
+
+    @JsonProperty
+    public String getRoutine()
+    {
+        return routine;
+    }
+
+    @JsonProperty
+    public String getAuthorization()
+    {
+        return authorization;
+    }
 }

presto-tests/src/test/java/io/prestosql/execution/TestEventListener.java

@@ -29,6 +29,7 @@
 import io.prestosql.spi.eventlistener.QueryCompletedEvent;
 import io.prestosql.spi.eventlistener.QueryCreatedEvent;
 import io.prestosql.spi.eventlistener.QueryFailureInfo;
+import io.prestosql.spi.eventlistener.RoutineInfo;
 import io.prestosql.spi.eventlistener.SplitCompletedEvent;
 import io.prestosql.spi.eventlistener.TableInfo;
 import io.prestosql.testing.DistributedQueryRunner;
@@ -219,7 +220,7 @@ public void testNormalQuery()
     }
 
     @Test
-    public void testReferencedTables()
+    public void testReferencedTablesAndRoutines()
             throws Exception
     {
         // We expect the following events
@@ -242,6 +243,13 @@ public void testReferencedTables()
         ColumnInfo column = table.getColumns().get(0);
         assertEquals(column.getColumn(), "linenumber");
         assertTrue(column.getMasks().isEmpty());
+
+        List<RoutineInfo> routines = event.getMetadata().getRoutines();
+        assertEquals(tables.size(), 1);
+
+        RoutineInfo routine = routines.get(0);
+        assertEquals(routine.getRoutine(), "sum");
+        assertEquals(routine.getAuthorization(), "user");
     }
 
     @Test