@Before
public void setupMocks() {
	context = Mockito.mock(Context.class);
	planner = Mockito.mock(RelOptPlanner.class);
	relTraitSet = RelTraitSet.createEmpty();
	superFactory = Mockito.mock(JdbcRelBuilderFactoryFactory.class);
	miniFactory = Mockito.mock(JdbcRelBuilderFactory.class);
	rule = Mockito.mock(ForcedRule.class);
	program = new ForcedRulesProgram(superFactory, rule);
	inNode = Mockito.mock(RelNode.class);
	relOptMaterializationList = Arrays.asList();
	relOptLatticeList = Arrays.asList();

	Mockito.doReturn(context).when(planner).getContext();
	Mockito.doReturn(miniFactory).when(superFactory).create(Mockito.same(context));
}
 

/**
 * Creates a configured {@link FlinkRelBuilder} for a planning session.
 *
 * @param currentCatalog the current default catalog to look for first during planning.
 * @param currentDatabase the current default database to look for first during planning.
 * @return configured rel builder
 */
public FlinkRelBuilder createRelBuilder(String currentCatalog, String currentDatabase) {
	RelOptCluster cluster = FlinkRelOptClusterFactory.create(
		planner,
		new RexBuilder(typeFactory));
	RelOptSchema relOptSchema = createCatalogReader(false, currentCatalog, currentDatabase);
	Context chain = Contexts.chain(
		context,
		// We need to overwrite the default scan factory, which does not
		// expand views. The expandingScanFactory uses the FlinkPlanner to translate a view
		// into a rel tree, before applying any subsequent rules.
		Contexts.of(RelFactories.expandingScanFactory(
			createFlinkPlanner(currentCatalog, currentDatabase),
			RelFactories.DEFAULT_TABLE_SCAN_FACTORY))
	);

	return new FlinkRelBuilder(chain, cluster, relOptSchema, expressionBridge);
}
 

protected RelBuilder(Context context, RelOptCluster cluster,
    RelOptSchema relOptSchema) {
  this.cluster = cluster;
  this.relOptSchema = relOptSchema;
  if (context == null) {
    context = Contexts.EMPTY_CONTEXT;
  }
  this.config = getConfig(context);
  this.viewExpander = getViewExpander(cluster, context);
  this.struct =
      Objects.requireNonNull(RelFactories.Struct.fromContext(context));
  final RexExecutor executor =
      Util.first(context.unwrap(RexExecutor.class),
          Util.first(cluster.getPlanner().getExecutor(), RexUtil.EXECUTOR));
  final RelOptPredicateList predicates = RelOptPredicateList.EMPTY;
  this.simplifier =
      new RexSimplify(cluster.getRexBuilder(), predicates, executor);
}
 

protected TesterImpl(DiffRepository diffRepos, boolean enableDecorrelate,
    boolean enableTrim, boolean enableExpand, boolean enableLateDecorrelate,
    boolean enableTypeCoercion,
    SqlTestFactory.MockCatalogReaderFactory catalogReaderFactory,
    Function<RelOptCluster, RelOptCluster> clusterFactory,
    SqlToRelConverter.Config config, SqlConformance conformance,
    Context context) {
  this.diffRepos = diffRepos;
  this.enableDecorrelate = enableDecorrelate;
  this.enableTrim = enableTrim;
  this.enableExpand = enableExpand;
  this.enableLateDecorrelate = enableLateDecorrelate;
  this.enableTypeCoercion = enableTypeCoercion;
  this.catalogReaderFactory = catalogReaderFactory;
  this.clusterFactory = clusterFactory;
  this.config = config;
  this.conformance = conformance;
  this.context = context;
}
 

static VolcanoPlanner createVolcanoPlanner() {
    RelOptCostFactory costFactory = RelOptCostImpl.FACTORY;
    Context externalContext = null;
    VolcanoPlanner volcanoPlanner = new VolcanoPlanner(costFactory, externalContext);
    // RexExecutor rexExecutor = null;
    return volcanoPlanner;
}
 

protected RelBuilder(Context context, RelOptCluster cluster, RelOptSchema relOptSchema) {
    this.cluster = cluster;
    this.relOptSchema = relOptSchema;
    if (context == null) {
        context = Contexts.EMPTY_CONTEXT;
    }
    this.simplify = Hook.REL_BUILDER_SIMPLIFY.get(true);
    this.aggregateFactory = Util.first(context.unwrap(RelFactories.AggregateFactory.class),
            RelFactories.DEFAULT_AGGREGATE_FACTORY);
    this.filterFactory = Util.first(context.unwrap(RelFactories.FilterFactory.class),
            RelFactories.DEFAULT_FILTER_FACTORY);
    this.projectFactory = Util.first(context.unwrap(RelFactories.ProjectFactory.class),
            RelFactories.DEFAULT_PROJECT_FACTORY);
    this.sortFactory = Util.first(context.unwrap(RelFactories.SortFactory.class),
            RelFactories.DEFAULT_SORT_FACTORY);
    this.exchangeFactory = Util.first(context.unwrap(RelFactories.ExchangeFactory.class),
            RelFactories.DEFAULT_EXCHANGE_FACTORY);
    this.sortExchangeFactory = Util.first(context.unwrap(RelFactories.SortExchangeFactory.class),
            RelFactories.DEFAULT_SORT_EXCHANGE_FACTORY);
    this.setOpFactory = Util.first(context.unwrap(RelFactories.SetOpFactory.class),
            RelFactories.DEFAULT_SET_OP_FACTORY);
    this.joinFactory = Util.first(context.unwrap(RelFactories.JoinFactory.class),
            RelFactories.DEFAULT_JOIN_FACTORY);
    this.semiJoinFactory = Util.first(context.unwrap(RelFactories.SemiJoinFactory.class),
            RelFactories.DEFAULT_SEMI_JOIN_FACTORY);
    this.correlateFactory = Util.first(context.unwrap(RelFactories.CorrelateFactory.class),
            RelFactories.DEFAULT_CORRELATE_FACTORY);
    this.valuesFactory = Util.first(context.unwrap(RelFactories.ValuesFactory.class),
            RelFactories.DEFAULT_VALUES_FACTORY);
    this.scanFactory = Util.first(context.unwrap(RelFactories.TableScanFactory.class),
            RelFactories.DEFAULT_TABLE_SCAN_FACTORY);
    this.snapshotFactory = Util.first(context.unwrap(RelFactories.SnapshotFactory.class),
            RelFactories.DEFAULT_SNAPSHOT_FACTORY);
    this.matchFactory = Util.first(context.unwrap(RelFactories.MatchFactory.class),
            RelFactories.DEFAULT_MATCH_FACTORY);
    final RexExecutor executor = Util.first(context.unwrap(RexExecutor.class),
            Util.first(cluster.getPlanner().getExecutor(), RexUtil.EXECUTOR));
    final RelOptPredicateList predicates = RelOptPredicateList.EMPTY;
    this.simplifier = new RexSimplify(cluster.getRexBuilder(), predicates, executor);
}
 

/**
 * Creates a {@code VolcanoPlanner} with a given cost factory.
 */
public VolcanoPlanner(RelOptCostFactory costFactory, //
    Context externalContext) {
  super(costFactory == null ? VolcanoCost.FACTORY : costFactory, //
      externalContext);
  this.zeroCost = this.costFactory.makeZeroCost();
}
 

/**
 * Creates a new HepPlanner with the option to keep the graph a
 * tree (noDag = true) or allow DAG (noDag = false).
 *
 * @param noDag      If false, create shared nodes if expressions are
 *                   identical
 * @param program    Program controlling rule application
 * @param onCopyHook Function to call when a node is copied
 */
public HepPlanner(
    HepProgram program,
    Context context,
    boolean noDag,
    Function2<RelNode, RelNode, Void> onCopyHook,
    RelOptCostFactory costFactory) {
  super(costFactory, context);
  this.mainProgram = program;
  this.onCopyHook = Util.first(onCopyHook, Functions.ignore2());
  this.noDag = noDag;
}
 

/** Creates a {@link RelBuilderFactory}, a partially-created DrillRelBuilder.
 * Just add a {@link RelOptCluster} and a {@link RelOptSchema} */
public static RelBuilderFactory proto(final Context context) {
  return new RelBuilderFactory() {
    public RelBuilder create(RelOptCluster cluster, RelOptSchema schema) {
      return new DrillRelBuilder(context, cluster, schema);
    }
  };
}
 

private RelDecorrelator(
    CorelMap cm,
    Context context,
    RelBuilder relBuilder) {
  this.cm = cm;
  this.context = context;
  this.relBuilder = relBuilder;
}
 

private RelDecorrelator(
    CorelMap cm,
    Context context,
    RelBuilder relBuilder) {
  this.cm = cm;
  this.context = context;
  this.relBuilder = relBuilder;
}
 

public DremioHepPlanner(final HepProgram program, final Context context, final RelOptCostFactory costFactory, PlannerPhase phase) {
  super(program, context, false, null, costFactory);
  this.cancelFlag = new CancelFlag(context.unwrap(PlannerSettings.class).getMaxPlanningPerPhaseMS(), TimeUnit.MILLISECONDS);
  this.phase = phase;
  this.listener = new MaxNodesListener(context.unwrap(PlannerSettings.class).getMaxNodesPerPlan());
  addListener(listener);
}
 

/** Creates a {@link RelBuilderFactory}, a partially-created RelBuilder.
 * Just add a {@link RelOptCluster} and a {@link RelOptSchema} */
public static RelBuilderFactory proto(final Context context) {
  return new RelBuilderFactory() {
    @Override
    public RelBuilder create(RelOptCluster cluster, RelOptSchema schema) {
      return new RelBuilder(context, cluster, schema);
    }
  };
}
 

protected DremioRelDecorrelator(
    CorelMap cm,
    Context context,
    RelBuilder relBuilder,
    boolean forceValueGenerator,
    boolean isRelPlanning) {
  super(cm, context, relBuilder, forceValueGenerator);
  this.isRelPlanning = isRelPlanning;
}
 

private DremioVolcanoPlanner(RelOptCostFactory costFactory, Context context, SubstitutionProvider substitutionProvider) {
  super(costFactory, context);
  this.substitutionProvider = substitutionProvider;
  this.cancelFlag = new CancelFlag(context.unwrap(PlannerSettings.class).getMaxPlanningPerPhaseMS(), TimeUnit.MILLISECONDS);
  this.phase = null;
  this.listener = new MaxNodesListener(context.unwrap(PlannerSettings.class).getMaxNodesPerPlan());
  addListener(listener);
}
 

public static DremioVolcanoPlanner of(RelOptCostFactory costFactory, Context context, SubstitutionProvider substitutionProvider, RexExecutor executor) {
  DremioVolcanoPlanner volcanoPlanner = new DremioVolcanoPlanner(costFactory, context, substitutionProvider);
  volcanoPlanner.setExecutor(executor);
  volcanoPlanner.clearRelTraitDefs();
  volcanoPlanner.addRelTraitDef(ConventionTraitDef.INSTANCE);
  volcanoPlanner.addRelTraitDef(DistributionTraitDef.INSTANCE);
  volcanoPlanner.addRelTraitDef(RelCollationTraitDef.INSTANCE);

  return volcanoPlanner;
}
 

/**
 * Creates a configured {@link FlinkRelBuilder} for a planning session.
 *
 * @param currentCatalog the current default catalog to look for first during planning.
 * @param currentDatabase the current default database to look for first during planning.
 * @return configured rel builder
 */
public FlinkRelBuilder createRelBuilder(String currentCatalog, String currentDatabase) {
	FlinkCalciteCatalogReader relOptSchema = createCatalogReader(
			false,
			currentCatalog,
			currentDatabase);

	Context chain = Contexts.of(
		context,
		// Sets up the ViewExpander explicitly for FlinkRelBuilder.
		createFlinkPlanner(currentCatalog, currentDatabase).createToRelContext()
	);
	return new FlinkRelBuilder(chain, cluster, relOptSchema);
}
 

public static @Nonnull Struct fromContext(Context context) {
  Struct struct = context.unwrap(Struct.class);
  if (struct != null) {
    return struct;
  }
  return new Struct(
      Util.first(context.unwrap(FilterFactory.class),
          DEFAULT_FILTER_FACTORY),
      Util.first(context.unwrap(ProjectFactory.class),
          DEFAULT_PROJECT_FACTORY),
      Util.first(context.unwrap(AggregateFactory.class),
          DEFAULT_AGGREGATE_FACTORY),
      Util.first(context.unwrap(SortFactory.class),
          DEFAULT_SORT_FACTORY),
      Util.first(context.unwrap(ExchangeFactory.class),
          DEFAULT_EXCHANGE_FACTORY),
      Util.first(context.unwrap(SortExchangeFactory.class),
          DEFAULT_SORT_EXCHANGE_FACTORY),
      Util.first(context.unwrap(SetOpFactory.class),
          DEFAULT_SET_OP_FACTORY),
      Util.first(context.unwrap(JoinFactory.class),
          DEFAULT_JOIN_FACTORY),
      Util.first(context.unwrap(CorrelateFactory.class),
          DEFAULT_CORRELATE_FACTORY),
      Util.first(context.unwrap(ValuesFactory.class),
          DEFAULT_VALUES_FACTORY),
      Util.first(context.unwrap(TableScanFactory.class),
          DEFAULT_TABLE_SCAN_FACTORY),
      Util.first(context.unwrap(TableFunctionScanFactory.class),
          DEFAULT_TABLE_FUNCTION_SCAN_FACTORY),
      Util.first(context.unwrap(SnapshotFactory.class),
          DEFAULT_SNAPSHOT_FACTORY),
      Util.first(context.unwrap(MatchFactory.class),
          DEFAULT_MATCH_FACTORY),
      Util.first(context.unwrap(SpoolFactory.class),
          DEFAULT_SPOOL_FACTORY),
      Util.first(context.unwrap(RepeatUnionFactory.class),
          DEFAULT_REPEAT_UNION_FACTORY));
}
 

protected RelDecorrelator(
    CorelMap cm,
    Context context,
    RelBuilder relBuilder) {
  this.cm = cm;
  this.context = context;
  this.relBuilder = relBuilder;
}
 

/** Derives the Config to be used for this RelBuilder.
 *
 * <p>Overrides {@link RelBuilder.Config#simplify} if
 * {@link Hook#REL_BUILDER_SIMPLIFY} is set.
 */
private Config getConfig(Context context) {
  final Config config =
      Util.first(context.unwrap(Config.class), Config.DEFAULT);
  boolean simplify = Hook.REL_BUILDER_SIMPLIFY.get(config.simplify());
  return config.withSimplify(simplify);
}
 

StdFrameworkConfig(Context context,
    SqlRexConvertletTable convertletTable,
    SqlOperatorTable operatorTable,
    ImmutableList<Program> programs,
    ImmutableList<RelTraitDef> traitDefs,
    SqlParser.Config parserConfig,
    SqlValidator.Config sqlValidatorConfig,
    SqlToRelConverter.Config sqlToRelConverterConfig,
    SchemaPlus defaultSchema,
    RelOptCostFactory costFactory,
    RelDataTypeSystem typeSystem,
    RexExecutor executor,
    boolean evolveLattice,
    SqlStatisticProvider statisticProvider,
    RelOptTable.ViewExpander viewExpander) {
  this.context = context;
  this.convertletTable = convertletTable;
  this.operatorTable = operatorTable;
  this.programs = programs;
  this.traitDefs = traitDefs;
  this.parserConfig = parserConfig;
  this.sqlValidatorConfig = sqlValidatorConfig;
  this.sqlToRelConverterConfig = sqlToRelConverterConfig;
  this.defaultSchema = defaultSchema;
  this.costFactory = costFactory;
  this.typeSystem = typeSystem;
  this.executor = executor;
  this.evolveLattice = evolveLattice;
  this.statisticProvider = statisticProvider;
  this.viewExpander = viewExpander;
}
 

/**
 * Creates a {@code VolcanoPlanner} with a given cost factory.
 */
public VolcanoPlanner(RelOptCostFactory costFactory,
    Context externalContext) {
  super(costFactory == null ? VolcanoCost.FACTORY : costFactory,
      externalContext);
  this.zeroCost = this.costFactory.makeZeroCost();
  // If LOGGER is debug enabled, enable provenance information to be captured
  this.provenanceMap = LOGGER.isDebugEnabled() ? new HashMap<>()
      : Util.blackholeMap();
}
 

/**
 * Creates a new HepPlanner with the option to keep the graph a
 * tree (noDag = true) or allow DAG (noDag = false).
 *
 * @param noDag      If false, create shared nodes if expressions are
 *                   identical
 * @param program    Program controlling rule application
 * @param onCopyHook Function to call when a node is copied
 */
public HepPlanner(
    HepProgram program,
    Context context,
    boolean noDag,
    Function2<RelNode, RelNode, Void> onCopyHook,
    RelOptCostFactory costFactory) {
  super(costFactory, context);
  this.mainProgram = program;
  this.onCopyHook = Util.first(onCopyHook, Functions.ignore2());
  this.noDag = noDag;
}
 

private RelDecorrelator(CorelMap cm, Context context, RelBuilder relBuilder) {
    this.cm = cm;
    this.context = context;
    this.relBuilder = relBuilder;
}
 

/** Creates a {@link RelBuilderFactory}, a partially-created RelBuilder.
 * Just add a {@link RelOptCluster} and a {@link RelOptSchema} */
public static RelBuilderFactory proto(final Context context) {
    return (cluster, schema) -> new RelBuilder(context, cluster, schema);
}
 

protected DrillRelBuilder(Context context, RelOptCluster cluster, RelOptSchema relOptSchema) {
  super(context, cluster, relOptSchema);
  this.filterFactory =
      Util.first(context.unwrap(RelFactories.FilterFactory.class),
          RelFactories.DEFAULT_FILTER_FACTORY);
}
 

public Context getContext() {
	return context;
}
 

private JdbcRelBuilder(Context context, RelOptCluster cluster, RelOptSchema relOptSchema) {
	super(context, cluster, relOptSchema);
}
 

public Factory(Context context) {
	this.context = context;
}
 

@Override
public JdbcRelBuilderFactory create(Context context) {
	return new JdbcRelBuilder.Factory(context);
}