Impose runtime limits on temp table and recursive queries

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/EvaluationContext.java

@@ -23,7 +23,9 @@
 import com.apple.foundationdb.annotation.API;
 import com.apple.foundationdb.record.logging.LogMessageKeys;
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
+import com.apple.foundationdb.record.query.plan.cascades.TempTable;
 import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository;
+import com.google.common.collect.ImmutableMap;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -45,7 +47,11 @@ public class EvaluationContext {
     @Nonnull
     private final TypeRepository typeRepository;
 
-    public static final EvaluationContext EMPTY = new EvaluationContext(Bindings.EMPTY_BINDINGS, TypeRepository.EMPTY_SCHEMA);
+    @Nonnull
+    private final TempTable.Factory tempTableFactory;
+
+    public static final EvaluationContext EMPTY = new EvaluationContext(Bindings.EMPTY_BINDINGS, TypeRepository.EMPTY_SCHEMA,
+            TempTable.Factory.instance());
 
     /**
      * Get an empty evaluation context.
@@ -56,9 +62,12 @@ public static EvaluationContext empty() {
         return EMPTY;
     }
 
-    private EvaluationContext(@Nonnull Bindings bindings, @Nonnull TypeRepository typeRepository) {
+    private EvaluationContext(@Nonnull final Bindings bindings,
+                              @Nonnull final TypeRepository typeRepository,
+                              @Nonnull final TempTable.Factory tempTableFactory) {
         this.bindings = bindings;
         this.typeRepository = typeRepository;
+        this.tempTableFactory = tempTableFactory;
     }
 
     /**
@@ -69,24 +78,65 @@ private EvaluationContext(@Nonnull Bindings bindings, @Nonnull TypeRepository ty
      */
     @Nonnull
     public static EvaluationContext forBindings(@Nonnull Bindings bindings) {
-        return new EvaluationContext(bindings, TypeRepository.EMPTY_SCHEMA);
+        return new EvaluationContext(bindings, TypeRepository.EMPTY_SCHEMA, TempTable.Factory.instance());
     }
 
     /**
-     * Create a new {@link EvaluationContext} around a given set of {@link Bindings} and a {@link TypeRepository}.
-     * from parameter names to values.
+     * Create a new {@link EvaluationContext} around a given set of {@link Bindings} from parameter names to values,
+     * and a {@link TypeRepository}.
      * @param bindings a mapping from parameter name to values
      * @param typeRepository a type repository
      * @return a new evaluation context with the bindings and the schema.
      */
     @Nonnull
-    public static EvaluationContext forBindingsAndTypeRepository(@Nonnull Bindings bindings, @Nonnull TypeRepository typeRepository) {
-        return new EvaluationContext(bindings, typeRepository);
+    public static EvaluationContext forBindingsAndTypeRepository(@Nonnull final Bindings bindings,
+                                                                 @Nonnull final TypeRepository typeRepository) {
+        return new EvaluationContext(bindings, typeRepository, TempTable.Factory.instance());
     }
 
     @Nonnull
-    public static EvaluationContext forTypeRepository(@Nonnull TypeRepository typeRepository) {
-        return new EvaluationContext(Bindings.EMPTY_BINDINGS, typeRepository);
+    public static EvaluationContext forTypeRepository(@Nonnull final TypeRepository typeRepository) {
+        return new EvaluationContext(Bindings.EMPTY_BINDINGS, typeRepository, TempTable.Factory.instance());
+    }
+
+    /**
+     * Create a new {@link EvaluationContext} around a given set of {@link Bindings} from parameter names to values,
+     * and a custom {@link TempTable.Factory}.
+     * @param bindings a mapping from parameter name to values
+     * @param tempTableFactory a factory of {@link TempTable} instances.
+     * @return a new evaluation context with the bindings and the schema.
+     */
+    @Nonnull
+    public static EvaluationContext forBindingsAndTempTableFactory(@Nonnull final Bindings bindings,
+                                                                   @Nonnull final TempTable.Factory tempTableFactory) {
+        return new EvaluationContext(bindings, TypeRepository.EMPTY_SCHEMA, tempTableFactory);
+    }
+
+    /**
+     * Create a new {@link EvaluationContext} around a custom {@link TempTable.Factory} and a {@link TypeRepository}.
+     * @param typeRepository a type repository
+     * @param tempTableFactory a factory of {@link TempTable} instances.
+     * @return a new evaluation context with the bindings and the schema.
+     */
+    @Nonnull
+    public static EvaluationContext forTypeRepositoryAndTempTableFactory(@Nonnull final TypeRepository typeRepository,
+                                                                         @Nonnull final TempTable.Factory tempTableFactory) {
+        return new EvaluationContext(Bindings.EMPTY_BINDINGS, typeRepository, tempTableFactory);
+    }
+
+    /**
+     * Create a new {@link EvaluationContext} around a given set of {@link Bindings} from parameter names to values,
+     * a {@link TypeRepository}, and a custom {@link TempTable.Factory}.
+     * @param bindings a mapping from parameter name to values
+     * @param typeRepository a type repository
+     * @param tempTableFactory a factory of {@link TempTable} instances.
+     * @return a new evaluation context with the bindings and the schema.
+     */
+    @Nonnull
+    public static EvaluationContext forBindingsAndTypeRepositoryAndTempTableFactory(@Nonnull final Bindings bindings,
+                                                                                    @Nonnull final TypeRepository typeRepository,
+                                                                                    @Nonnull final TempTable.Factory tempTableFactory) {
+        return new EvaluationContext(bindings, typeRepository, TempTable.Factory.instance());
     }
 
     /**
@@ -98,7 +148,8 @@ public static EvaluationContext forTypeRepository(@Nonnull TypeRepository typeRe
      */
     @Nonnull
     public static EvaluationContext forBinding(@Nonnull String bindingName, @Nullable Object value) {
-        return new EvaluationContext(Bindings.newBuilder().set(bindingName, value).build(), TypeRepository.EMPTY_SCHEMA);
+        return new EvaluationContext(Bindings.newBuilder().set(bindingName, value).build(), TypeRepository.EMPTY_SCHEMA,
+                TempTable.Factory.instance());
     }
 
     /**
@@ -209,7 +260,29 @@ public EvaluationContext withBinding(@Nonnull String bindingName, @Nullable Obje
      *
      * @return a new <code>EvaluationContext</code> with the new binding
      */
-    public EvaluationContext withBinding(final Bindings.Internal type, @Nonnull CorrelationIdentifier alias, @Nullable Object value) {
+    @Nonnull
+    public EvaluationContext withBinding(@Nonnull final Bindings.Internal type,
+                                         @Nonnull final CorrelationIdentifier alias,
+                                         @Nullable final Object value) {
         return childBuilder().setBinding(type.bindingName(alias.getId()), value).build(typeRepository);
     }
+
+    @Nonnull
+    public EvaluationContext withNewTempTableBinding(@Nonnull final Bindings.Internal type,
+                                                     @Nonnull final CorrelationIdentifier alias) {
+        final var tempTable = getTempTableFactory().createTempTable();
+        // CONSTANT and CORRELATION types expect a different layout of the values.
+        // ideally this should be streamlined differently.
+        if (type == Bindings.Internal.CONSTANT) {
+            final ImmutableMap.Builder<String, Object> constants = ImmutableMap.builder();
+            constants.put(alias.getId(), tempTable);
+            return childBuilder().setBinding(type.bindingName(alias.getId()), constants.build()).build(typeRepository);
+        }
+        return childBuilder().setBinding(type.bindingName(alias.getId()), getTempTableFactory().createTempTable()).build(typeRepository);
+    }
+
+    @Nonnull
+    public TempTable.Factory getTempTableFactory() {
+        return tempTableFactory;
+    }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/EvaluationContextBuilder.java

@@ -22,6 +22,7 @@
 
 import com.apple.foundationdb.annotation.API;
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
+import com.apple.foundationdb.record.query.plan.cascades.TempTable;
 import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository;
 
 import javax.annotation.Nonnull;
@@ -37,12 +38,15 @@
 public class EvaluationContextBuilder {
     @Nonnull
     protected final Bindings.Builder bindings;
+    @Nonnull
+    protected final TempTable.Factory tempTableFactory;
 
     /**
      * Create an empty builder.
      */
     protected EvaluationContextBuilder() {
         this.bindings = Bindings.newBuilder();
+        this.tempTableFactory = TempTable.Factory.instance();
     }
 
     /**
@@ -54,6 +58,7 @@ protected EvaluationContextBuilder() {
      */
     protected EvaluationContextBuilder(@Nonnull EvaluationContext original) {
         this.bindings = original.getBindings().childBuilder();
+        this.tempTableFactory = original.getTempTableFactory();
     }
 
     /**
@@ -116,6 +121,6 @@ public EvaluationContextBuilder setConstant(@Nonnull CorrelationIdentifier alias
      */
     @Nonnull
     public EvaluationContext build(@Nonnull final TypeRepository typeRepository) {
-        return EvaluationContext.forBindingsAndTypeRepository(bindings.build(), typeRepository);
+        return EvaluationContext.forBindingsAndTypeRepositoryAndTempTableFactory(bindings.build(), typeRepository, tempTableFactory);
     }
 }