Make AuthContext available to QPL, QOPL, and UPL (#4115)

* Record AuthContext when creating QueryPerformanceNugget (for QPL and QOPL) and PerformanceEntry (for UPL). Add AuthContext (toString()ed) as a column in QPL, QOPL, and UPL (for the in-memory blink tables available in core). Enables registered loggers to also get AuthContext. * Add in the enclosing AuthContext in AbstractScriptSession.evaluateScript.
deephaven · Jun 30, 2023 · 709f2a4 · 709f2a4
1 parent b497fea
commit 709f2a4
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 10 deletions.
diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/PerformanceEntry.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/PerformanceEntry.java
@@ -3,7 +3,9 @@
  */
 package io.deephaven.engine.table.impl.perf;
 
+import io.deephaven.auth.AuthContext;
 import io.deephaven.base.log.LogOutput;
+import io.deephaven.engine.context.ExecutionContext;
 import io.deephaven.engine.rowset.RowSet;
 import io.deephaven.engine.rowset.RowSetShiftData;
 import io.deephaven.engine.table.TableListener;
@@ -22,6 +24,8 @@ public class PerformanceEntry extends BasePerformanceEntry implements TableListe
     private final String description;
     private final String callerLine;
 
+    private final AuthContext authContext;
+
     private long intervalInvocationCount;
 
     private long intervalAdded;
@@ -44,6 +48,7 @@ public class PerformanceEntry extends BasePerformanceEntry implements TableListe
         this.operationNumber = operationNumber;
         this.description = description;
         this.callerLine = callerLine;
+        authContext = id == QueryConstants.NULL_INT ? null : ExecutionContext.getContext().getAuthContext();
         startSample = new RuntimeMemory.Sample();
         endSample = new RuntimeMemory.Sample();
         maxTotalMemory = 0;
@@ -114,6 +119,7 @@ public LogOutput append(final LogOutput logOutput) {
                 .append(", operationNumber=").append(operationNumber)
                 .append(", description='").append(description).append('\'')
                 .append(", callerLine='").append(callerLine).append('\'')
+                .append(", authContext=").append(authContext)
                 .append(", intervalUsageNanos=").append(getIntervalUsageNanos())
                 .append(", intervalCpuNanos=").append(getIntervalCpuNanos())
                 .append(", intervalUserCpuNanos=").append(getIntervalUserCpuNanos())
@@ -152,6 +158,13 @@ public String getCallerLine() {
         return callerLine;
     }
 
+    /**
+     * @return The {@link AuthContext} that was installed when this PerformanceEntry was constructed
+     */
+    public AuthContext getAuthContext() {
+        return authContext;
+    }
+
     public long getIntervalAdded() {
         return intervalAdded;
     }

diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java
@@ -3,6 +3,8 @@
  */
 package io.deephaven.engine.table.impl.perf;
 
+import io.deephaven.auth.AuthContext;
+import io.deephaven.engine.context.ExecutionContext;
 import io.deephaven.time.DateTimeUtils;
 import io.deephaven.engine.table.impl.util.RuntimeMemory;
 import io.deephaven.util.QueryConstants;
@@ -35,9 +37,11 @@ public class QueryPerformanceNugget implements Serializable, AutoCloseable {
     private final int depth;
     private final String description;
     private final boolean isUser;
-    private final String callerLine;
     private final long inputSize;
 
+    private final AuthContext authContext;
+    private final String callerLine;
+
     private final long startClockTime;
 
     private final long startTimeNanos;
@@ -98,14 +102,15 @@ public class QueryPerformanceNugget implements Serializable, AutoCloseable {
         this.isUser = isUser;
         this.inputSize = inputSize;
 
+        authContext = ExecutionContext.getContext().getAuthContext();
+        callerLine = QueryPerformanceRecorder.getCallerLine();
+
         final RuntimeMemory runtimeMemory = RuntimeMemory.getInstance();
         runtimeMemory.read(startMemorySample);
 
         startAllocatedBytes = ThreadProfiler.DEFAULT.getCurrentThreadAllocatedBytes();
         startPoolAllocatedBytes = QueryPerformanceRecorder.getPoolAllocatedBytesForCurrentThread();
 
-        callerLine = QueryPerformanceRecorder.getCallerLine();
-
         startClockTime = System.currentTimeMillis();
         startTimeNanos = System.nanoTime();
 
@@ -128,11 +133,12 @@ private QueryPerformanceNugget() {
         isUser = false;
         inputSize = NULL_LONG;
 
+        authContext = null;
+        callerLine = null;
+
         startAllocatedBytes = NULL_LONG;
         startPoolAllocatedBytes = NULL_LONG;
 
-        callerLine = null;
-
         startClockTime = NULL_LONG;
         startTimeNanos = NULL_LONG;
 
@@ -251,6 +257,13 @@ public long getInputSize() {
         return inputSize;
     }
 
+    /**
+     * @return The {@link AuthContext} that was installed when this QueryPerformanceNugget was constructed
+     */
+    public AuthContext getAuthContext() {
+        return authContext;
+    }
+
     public String getCallerLine() {
         return callerLine;
     }

diff --git a/...e/src/main/java/io/deephaven/engine/table/impl/perf/UpdatePerformanceStreamPublisher.java b/...e/src/main/java/io/deephaven/engine/table/impl/perf/UpdatePerformanceStreamPublisher.java
@@ -42,7 +42,8 @@ class UpdatePerformanceStreamPublisher implements StreamPublisher {
             ColumnDefinition.ofLong("Collections"),
             ColumnDefinition.ofLong("CollectionTimeNanos"),
             ColumnDefinition.ofLong("EntryIntervalAllocatedBytes"),
-            ColumnDefinition.ofLong("EntryIntervalPoolAllocatedBytes"));
+            ColumnDefinition.ofLong("EntryIntervalPoolAllocatedBytes"),
+            ColumnDefinition.ofString("AuthContext"));
 
     public static TableDefinition definition() {
         return DEFINITION;
@@ -91,6 +92,7 @@ public synchronized void add(IntervalLevelDetails intervalLevelDetails, Performa
         chunks[20].asWritableLongChunk().add(performanceEntry.getCollectionTimeNanos());
         chunks[21].asWritableLongChunk().add(performanceEntry.getIntervalAllocatedBytes());
         chunks[22].asWritableLongChunk().add(performanceEntry.getIntervalPoolAllocatedBytes());
+        chunks[23].<String>asWritableObjectChunk().add(Objects.toString(performanceEntry.getAuthContext()));
         if (chunks[0].size() == CHUNK_SIZE) {
             flushInternal();
         }

diff --git a/...in/java/io/deephaven/engine/table/impl/util/QueryOperationPerformanceStreamPublisher.java b/...in/java/io/deephaven/engine/table/impl/util/QueryOperationPerformanceStreamPublisher.java
@@ -42,7 +42,8 @@ class QueryOperationPerformanceStreamPublisher implements StreamPublisher {
             ColumnDefinition.ofLong("AllocatedBytes"),
             ColumnDefinition.ofLong("PoolAllocatedBytes"),
             ColumnDefinition.ofLong("InputSizeLong"),
-            ColumnDefinition.ofBoolean("WasInterrupted"));
+            ColumnDefinition.ofBoolean("WasInterrupted"),
+            ColumnDefinition.ofString("AuthContext"));
     private static final int CHUNK_SIZE = ArrayBackedColumnSource.BLOCK_SIZE;
 
     public static TableDefinition definition() {
@@ -93,6 +94,7 @@ public synchronized void add(
         chunks[18].asWritableLongChunk().add(nugget.getAllocatedBytes());
         chunks[19].asWritableLongChunk().add(nugget.getPoolAllocatedBytes());
         chunks[20].asWritableByteChunk().add(BooleanUtils.booleanAsByte(nugget.wasInterrupted()));
+        chunks[21].<String>asWritableObjectChunk().add(Objects.toString(nugget.getAuthContext()));
         if (chunks[0].size() == CHUNK_SIZE) {
             flushInternal();
         }

diff --git a/...le/src/main/java/io/deephaven/engine/table/impl/util/QueryPerformanceStreamPublisher.java b/...le/src/main/java/io/deephaven/engine/table/impl/util/QueryPerformanceStreamPublisher.java
@@ -40,7 +40,8 @@ class QueryPerformanceStreamPublisher implements StreamPublisher {
             ColumnDefinition.ofLong("PoolAllocatedBytes"),
             ColumnDefinition.ofBoolean("WasInterrupted"),
             ColumnDefinition.ofBoolean("IsReplayer"),
-            ColumnDefinition.ofString("Exception"));
+            ColumnDefinition.ofString("Exception"),
+            ColumnDefinition.ofString("AuthContext"));
     private static final int CHUNK_SIZE = ArrayBackedColumnSource.BLOCK_SIZE;
 
     public static TableDefinition definition() {
@@ -124,6 +125,9 @@ public synchronized void add(
         // ColumnDefinition.ofString("Exception")
         chunks[17].<String>asWritableObjectChunk().add(queryProcessingResults.getException());
 
+        // ColumnDefinition.ofString("AuthContext")
+        chunks[18].<String>asWritableObjectChunk().add(Objects.toString(nugget.getAuthContext()));
+
         if (chunks[0].size() == CHUNK_SIZE) {
             flushInternal();
         }

diff --git a/engine/table/src/main/java/io/deephaven/engine/util/AbstractScriptSession.java b/engine/table/src/main/java/io/deephaven/engine/util/AbstractScriptSession.java
@@ -146,8 +146,10 @@ public synchronized final Changes evaluateScript(final String script, final @Nul
                 final SafeCloseable ignored = LivenessScopeStack.open(this, false)) {
 
             try {
-                // actually evaluate the script
-                executionContext.apply(() -> evaluate(script, scriptName));
+                // Actually evaluate the script; use the enclosing auth context, since AbstractScriptSession's
+                // ExecutionContext never has a non-null AuthContext
+                executionContext.withAuthContext(ExecutionContext.getContext().getAuthContext())
+                        .apply(() -> evaluate(script, scriptName));
             } catch (final RuntimeException err) {
                 evaluateErr = err;
             }