Support result in LoggingInvocationHandler

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/LoggingInvocationHandler.java

@@ -42,12 +42,19 @@ public class LoggingInvocationHandler
     private final Object delegate;
     private final ParameterNamesProvider parameterNames;
     private final Consumer<String> logger;
+    private final boolean includeResult;
 
     public LoggingInvocationHandler(Object delegate, ParameterNamesProvider parameterNames, Consumer<String> logger)
+    {
+        this(delegate, parameterNames, logger, false);
+    }
+
+    public LoggingInvocationHandler(Object delegate, ParameterNamesProvider parameterNames, Consumer<String> logger, boolean includeResult)
     {
         this.delegate = requireNonNull(delegate, "delegate is null");
         this.parameterNames = requireNonNull(parameterNames, "parameterNames is null");
         this.logger = requireNonNull(logger, "logger is null");
+        this.includeResult = includeResult;
     }
 
     @Override
@@ -66,7 +73,12 @@ protected Object handleInvocation(Object proxy, Method method, Object[] args)
             throw t;
         }
         Duration elapsed = Duration.nanosSince(startNanos);
-        logger.accept(format("%s succeeded in %s", invocationDescription(method, args), elapsed));
+        if (includeResult) {
+            logger.accept(format("%s succeeded in %s and returned %s", invocationDescription(method, args), elapsed, formatArgument(result)));
+        }
+        else {
+            logger.accept(format("%s succeeded in %s", invocationDescription(method, args), elapsed));
+        }
         return result;
     }
 

lib/trino-plugin-toolkit/src/test/java/io/trino/plugin/base/util/TestLoggingInvocationHandler.java

@@ -34,11 +34,12 @@ public void testLoggingAndExceptions()
         SomeInterface delegate = new SomeInterface()
         {
             @Override
-            public void run(boolean ok, String s)
+            public String run(boolean ok, String s)
             {
                 if (!ok) {
                     throw new ArrayStoreException(s);
                 }
+                return null;
             }
         };
         List<String> messages = new ArrayList<>();
@@ -60,8 +61,44 @@ public void run(boolean ok, String s)
                 });
     }
 
+    @Test
+    public void testLoggingResult()
+    {
+        SomeInterface delegate = new SomeInterface()
+        {
+            @Override
+            public String run(boolean ok, String s)
+            {
+                if (!ok) {
+                    throw new ArrayStoreException(s);
+                }
+                return "result=" + s;
+            }
+        };
+        List<String> messages = new ArrayList<>();
+        InvocationHandler handler = new LoggingInvocationHandler(delegate, new ReflectiveParameterNamesProvider(), messages::add, true);
+        SomeInterface proxy = newProxy(SomeInterface.class, handler);
+
+        proxy.run(true, "xyz");
+
+        assertThatThrownBy(() -> proxy.run(false, "bad"))
+                .isInstanceOf(ArrayStoreException.class)
+                .hasMessage("bad");
+
+        assertThat(messages)
+                .hasSize(2)
+                .satisfies(list -> {
+                    assertThat(list.get(0)).matches("\\QInvocation of run(ok=true, s='xyz') succeeded in\\E " + DURATION_PATTERN + " and returned 'result=xyz'");
+                    assertThat(list.get(1)).matches("\\QInvocation of run(ok=false, s='bad') took\\E " + DURATION_PATTERN +
+                            " \\Qand failed with java.lang.ArrayStoreException: bad\\E");
+                });
+    }
+
     private interface SomeInterface
     {
-        default void run(boolean ok, String s) {}
+        default String run(boolean ok, String s)
+        {
+            return null;
+        }
     }
 }