Extract variable api from ScriptSession, let ScriptSession guard reads

Integrations/src/main/java/io/deephaven/integrations/python/PythonDeephavenSession.java

@@ -24,14 +24,14 @@
 import io.deephaven.plugin.type.ObjectTypeLookup.NoOp;
 import io.deephaven.util.SafeCloseable;
 import io.deephaven.util.annotations.ScriptApi;
-import io.deephaven.util.annotations.VisibleForTesting;
 import io.deephaven.util.thread.NamingThreadFactory;
 import io.deephaven.util.thread.ThreadInitializationFactory;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jpy.KeyError;
 import org.jpy.PyDictWrapper;
 import org.jpy.PyInputMode;
+import org.jpy.PyLib;
 import org.jpy.PyLib.CallableKind;
 import org.jpy.PyModule;
 import org.jpy.PyObject;
@@ -41,14 +41,13 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.Collections;
-import java.util.LinkedHashMap;
-import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.function.Consumer;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 /**
@@ -150,13 +149,6 @@ public Thread newThread(@NotNull Runnable r) {
         }
     }
 
-    @Override
-    @VisibleForTesting
-    public QueryScope newQueryScope() {
-        // depend on the GIL instead of local synchronization
-        return new UnsynchronizedScriptSessionQueryScope(this);
-    }
-
     /**
      * Finds the specified script; and runs it as a file, or if it is a stream writes it to a temporary file in order to
      * run it.
@@ -183,20 +175,14 @@ private void runScript(String script) throws IOException {
         }
     }
 
-    @NotNull
+    @SuppressWarnings("unchecked")
     @Override
-    public Object getVariable(String name) throws QueryScope.MissingVariableException {
-        return scope
+    protected <T> T getVariable(String name) throws QueryScope.MissingVariableException {
+        return (T) scope
                 .getValue(name)
-                .orElseThrow(() -> new QueryScope.MissingVariableException("No variable for: " + name));
+                .orElseThrow(() -> new QueryScope.MissingVariableException("Missing variable " + name));
     }
 
-    @Override
-    public <T> T getVariable(String name, T defaultValue) {
-        return scope
-                .<T>getValueUnchecked(name)
-                .orElse(defaultValue);
-    }
 
     @SuppressWarnings("unused")
     @ScriptApi
@@ -224,13 +210,6 @@ protected void evaluate(String command, String scriptName) {
         }
     }
 
-    @Override
-    public Map<String, Object> getVariables() {
-        final Map<String, Object> outMap = new LinkedHashMap<>();
-        scope.getEntriesMap().forEach((key, value) -> outMap.put(key, maybeUnwrap(value)));
-        return outMap;
-    }
-
     protected static class PythonSnapshot implements Snapshot, SafeCloseable {
 
         private final PyDictWrapper dict;
@@ -295,34 +274,49 @@ private Object maybeUnwrap(PyObject o) {
     }
 
     @Override
-    public Set<String> getVariableNames() {
-        return Collections.unmodifiableSet(scope.getKeys().collect(Collectors.toSet()));
+    protected Set<String> getVariableNames(Predicate<String> allowName) {
+        return PyLib.ensureGil(() -> scope.getKeys().filter(allowName).collect(Collectors.toUnmodifiableSet()));
     }
 
     @Override
-    public boolean hasVariableName(String name) {
+    protected boolean hasVariable(String name) {
         return scope.containsKey(name);
     }
 
     @Override
-    public synchronized void setVariable(String name, @Nullable Object newValue) {
-        final PyDictWrapper globals = scope.mainGlobals();
-        if (newValue == null) {
-            try {
-                globals.delItem(name);
-            } catch (KeyError key) {
-                // ignore
-            }
-        } else {
-            if (!(newValue instanceof PyObject)) {
-                newValue = PythonObjectWrapper.wrap(newValue);
+    protected synchronized Object setVariable(String name, @Nullable Object newValue) {
+        Object old = PyLib.ensureGil(() -> {
+            final PyDictWrapper globals = scope.mainGlobals();
+
+            if (newValue == null) {
+                try {
+                    return globals.unwrap().callMethod("pop", name);
+                } catch (KeyError key) {
+                    return null;
+                }
+            } else {
+                Object wrapped;
+                if (newValue instanceof PyObject) {
+                    wrapped = newValue;
+                } else {
+                    wrapped = PythonObjectWrapper.wrap(newValue);
+                }
+                // This isn't thread safe, we're relying on the GIL being kind to us (as we have historically done).
+                // There
+                // is no built-in for "replace a variable and return the old one".
+                Object prev = globals.get(name);
+                globals.setItem(name, wrapped);
+                return prev;
             }
-            globals.setItem(name, newValue);
-        }
+        });
 
         // Observe changes from this "setVariable" (potentially capturing previous or concurrent external changes from
         // other threads)
         observeScopeChanges();
+
+        // This doesn't return the same Java instance of PyObject, so we won't decref it properly, but
+        // again, that is consistent with how we've historically treated these references.
+        return old;
     }
 
     @Override

engine/context/src/main/java/io/deephaven/engine/context/EmptyQueryScope.java

@@ -40,9 +40,4 @@ public <T> T readParamValue(String name, T defaultValue) {
     public <T> void putParam(String name, T value) {
         throw new IllegalStateException("EmptyQueryScope cannot create parameters");
     }
-
-    @Override
-    public void putObjectFields(Object object) {
-        throw new IllegalStateException("EmptyQueryScope cannot create parameters");
-    }
 }

engine/context/src/main/java/io/deephaven/engine/context/ExecutionContext.java

@@ -172,6 +172,22 @@ public ExecutionContext withOperationInitializer(final OperationInitializer oper
                 operationInitializer);
     }
 
+    /**
+     * Returns, or creates, an execution context with the given value for {@code queryScope} and existing values for the
+     * other members.
+     *
+     * @param queryScope the query scope to use instead
+     * @return the execution context
+     */
+    public ExecutionContext withQueryScope(QueryScope queryScope) {
+        if (queryScope == this.queryScope) {
+            return this;
+        }
+        return new ExecutionContext(isSystemic, authContext, queryLibrary, queryScope, queryCompiler, updateGraph,
+                operationInitializer);
+    }
+
+
     /**
      * Execute runnable within this execution context.
      */

engine/context/src/main/java/io/deephaven/engine/context/PoisonedQueryScope.java

@@ -46,9 +46,4 @@ public <T> T readParamValue(String name, T defaultValue) {
     public <T> void putParam(String name, T value) {
         fail();
     }
-
-    @Override
-    public void putObjectFields(Object object) {
-        fail();
-    }
 }

engine/context/src/main/java/io/deephaven/engine/context/QueryScope.java

@@ -3,6 +3,10 @@
  */
 package io.deephaven.engine.context;
 
+import io.deephaven.engine.liveness.Liveness;
+import io.deephaven.engine.liveness.LivenessReferent;
+import io.deephaven.engine.liveness.ReferenceCountedLivenessNode;
+import io.deephaven.engine.updategraph.DynamicNode;
 import io.deephaven.time.DateTimeUtils;
 import io.deephaven.hash.KeyedObjectHashMap;
 import io.deephaven.hash.KeyedObjectKey;
@@ -21,7 +25,7 @@
 /**
  * Variable scope used to resolve parameter values during query execution.
  */
-public abstract class QueryScope implements LogOutputAppendable {
+public abstract class QueryScope extends ReferenceCountedLivenessNode implements LogOutputAppendable {
 
     /**
      * Adds a parameter to the default instance {@link QueryScope}, or updates the value of an existing parameter.
@@ -34,15 +38,6 @@ public static <T> void addParam(final String name, final T value) {
         ExecutionContext.getContext().getQueryScope().putParam(name, value);
     }
 
-    /**
-     * Adds an object's declared fields to the scope.
-     *
-     * @param object object whose fields will be added.
-     */
-    public static void addObjectFields(final Object object) {
-        ExecutionContext.getContext().getQueryScope().putObjectFields(object);
-    }
-
     /**
      * Gets a parameter from the default instance {@link QueryScope}.
      *
@@ -119,6 +114,12 @@ private static Object applyValueConversions(final Object value) {
         return value;
     }
 
+    protected QueryScope() {
+        super(false);
+
+        retainReference();
+    }
+
     // -----------------------------------------------------------------------------------------------------------------
     // Scope manipulation helper methods
     // -----------------------------------------------------------------------------------------------------------------
@@ -194,12 +195,15 @@ public final QueryScopeParam[] getParams(final Collection<String> names) throws
     public abstract <T> void putParam(final String name, final T value);
 
     /**
-     * Add an object's public members (referenced reflectively, not a shallow copy!) to this scope if supported.
-     * <b>Note:</b> This is an optional method.
+     * Asks the session to remove any wrapping that exists on scoped objects so that clients can fetch them. Defaults to
+     * returning the object itself.
      *
-     * @param object object to add public members from.
+     * @param object the scoped object
+     * @return an obj which can be consumed by a client
      */
-    public abstract void putObjectFields(final Object object);
+    public Object unwrapObject(Object object) {
+        return object;
+    }
 
     // -----------------------------------------------------------------------------------------------------------------
     // LogOutputAppendable implementation
@@ -232,7 +236,11 @@ public static class StandaloneImpl extends QueryScope {
         private final KeyedObjectHashMap<String, ValueRetriever> valueRetrievers =
                 new KeyedObjectHashMap<>(new ValueRetrieverNameKey());
 
-        public StandaloneImpl() {}
+        @Override
+        protected void destroy() {
+            super.destroy();
+            valueRetrievers.clear();
+        }
 
         @Override
         public Set<String> getParamNames() {
@@ -276,14 +284,19 @@ public <T> T readParamValue(final String name, final T defaultValue) {
 
         @Override
         public <T> void putParam(final String name, final T value) {
+            if (value instanceof LivenessReferent && DynamicNode.notDynamicOrIsRefreshing(value)) {
+                manage((LivenessReferent) value);
+            }
             NameValidator.validateQueryParameterName(name);
             // TODO: Can I get rid of this applyValueConversions? It's too inconsistent to feel safe.
-            valueRetrievers.put(name, new SimpleValueRetriever<>(name, applyValueConversions(value)));
-        }
+            ValueRetriever<?> oldValueRetriever =
+                    valueRetrievers.put(name, new SimpleValueRetriever<>(name, applyValueConversions(value)));
 
-        public void putObjectFields(final Object object) {
-            for (final Field field : object.getClass().getDeclaredFields()) {
-                valueRetrievers.put(field.getName(), new ReflectiveValueRetriever(object, field));
+            if (oldValueRetriever != null) {
+                Object oldValue = oldValueRetriever.getValue();
+                if (oldValue instanceof LivenessReferent && DynamicNode.notDynamicOrIsRefreshing(oldValue)) {
+                    unmanage((LivenessReferent) oldValue);
+                }
             }
         }
 

engine/sql/src/main/java/io/deephaven/engine/sql/Sql.java

@@ -3,12 +3,11 @@
 import io.deephaven.base.log.LogOutput;
 import io.deephaven.base.log.LogOutput.ObjFormatter;
 import io.deephaven.engine.context.ExecutionContext;
+import io.deephaven.engine.context.QueryScope;
 import io.deephaven.engine.table.ColumnDefinition;
 import io.deephaven.engine.table.Table;
 import io.deephaven.engine.table.TableDefinition;
 import io.deephaven.engine.table.impl.TableCreatorImpl;
-import io.deephaven.engine.util.AbstractScriptSession.ScriptSessionQueryScope;
-import io.deephaven.engine.util.ScriptSession;
 import io.deephaven.internal.log.LoggerFactory;
 import io.deephaven.io.logger.Logger;
 import io.deephaven.qst.column.header.ColumnHeader;
@@ -83,18 +82,16 @@ private static Map<String, Table> currentScriptSessionNamedTables() {
         final Map<String, Table> scope = new HashMap<>();
         // getVariables() is inefficient
         // See SQLTODO(catalog-reader-implementation)
-        for (Entry<String, Object> e : currentScriptSession().getVariables().entrySet()) {
-            if (e.getValue() instanceof Table) {
-                scope.put(e.getKey(), (Table) e.getValue());
+        QueryScope queryScope = ExecutionContext.getContext().getQueryScope();
+        for (String name : queryScope.getParamNames()) {
+            Object paramValue = queryScope.unwrapObject(queryScope.readParamValue(name));
+            if (paramValue instanceof Table) {
+                scope.put(name, (Table) paramValue);
             }
         }
         return scope;
     }
 
-    private static ScriptSession currentScriptSession() {
-        return ((ScriptSessionQueryScope) ExecutionContext.getContext().getQueryScope()).scriptSession();
-    }
-
     private static TableHeader adapt(TableDefinition tableDef) {
         final Builder builder = TableHeader.builder();
         for (ColumnDefinition<?> cd : tableDef.getColumns()) {