content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to you under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.calcite.rel.core;
  
  import org.apache.calcite.plan.RelOptCluster;
  import org.apache.calcite.plan.RelOptCost;
  import org.apache.calcite.plan.RelOptPlanner;
  import org.apache.calcite.plan.RelOptSchema;
  import org.apache.calcite.plan.RelOptTable;
  import org.apache.calcite.plan.RelOptUtil;
  import org.apache.calcite.plan.RelTraitSet;
  import org.apache.calcite.prepare.Prepare;
  import org.apache.calcite.rel.RelInput;
  import org.apache.calcite.rel.RelNode;
  import org.apache.calcite.rel.RelWriter;
  import org.apache.calcite.rel.SingleRel;
  import org.apache.calcite.rel.externalize.RelEnumTypes;
  import org.apache.calcite.rel.metadata.RelMetadataQuery;
  import org.apache.calcite.rel.type.RelDataType;
  import org.apache.calcite.rel.type.RelDataTypeFactory;
  import org.apache.calcite.rex.RexNode;
  import org.apache.calcite.sql.SqlKind;
  import org.apache.calcite.sql.type.SqlTypeUtil;
  
  import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  import java.util.List;
  
  import static com.google.common.base.Preconditions.checkArgument;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Relational expression that modifies a table.
   *
   * <p>It is similar to {@link org.apache.calcite.rel.core.TableScan},
   * but represents a request to modify a table rather than read from it.
   * It takes one child which produces the modified rows. Those rows are:
   *
   * <ul>
   * <li>For {@code INSERT}, those rows are the new values;
   * <li>for {@code DELETE}, the old values;
   * <li>for {@code UPDATE}, all old values plus updated new values.
   * </ul>
   */
  public abstract class TableModify extends SingleRel {
    //~ Enums ------------------------------------------------------------------
  
    /**
     * Enumeration of supported modification operations.
     */
    public enum Operation {
      INSERT, UPDATE, DELETE, MERGE
    }
  
    //~ Instance fields --------------------------------------------------------
  
    /**
     * The connection to the optimizing session.
     */
    protected final Prepare.CatalogReader catalogReader;
  
    /**
     * The table definition.
     */
    protected final RelOptTable table;
    private final Operation operation;
    private final @Nullable List<String> updateColumnList;
    private final @Nullable List<RexNode> sourceExpressionList;
    private @MonotonicNonNull RelDataType inputRowType;
    private final boolean flattened;
  
    //~ Constructors -----------------------------------------------------------
  
    /**
     * Creates a {@code TableModify}.
     *
     * <p>The UPDATE operation has format like this:
     * <blockquote>
     *   <pre>UPDATE table SET iden1 = exp1, ident2 = exp2  WHERE condition</pre>
     * </blockquote>
     *
     * @param cluster    Cluster this relational expression belongs to
     * @param traitSet   Traits of this relational expression
     * @param table      Target table to modify
     * @param catalogReader accessor to the table metadata.
     * @param input      Sub-query or filter condition
     * @param operation  Modify operation (INSERT, UPDATE, DELETE)
     * @param updateColumnList List of column identifiers to be updated
     *           (e.g. ident1, ident2); null if not UPDATE
     * @param sourceExpressionList List of value expressions to be set
     *           (e.g. exp1, exp2); null if not UPDATE
     * @param flattened Whether set flattens the input row type
     */
    protected TableModify(
        RelOptCluster cluster,
        RelTraitSet traitSet,
        RelOptTable table,
        Prepare.CatalogReader catalogReader,
        RelNode input,
        Operation operation,
        @Nullable List<String> updateColumnList,
        @Nullable List<RexNode> sourceExpressionList,
        boolean flattened) {
      super(cluster, traitSet, input);
      this.table = table;
      this.catalogReader = catalogReader;
      this.operation = operation;
      this.updateColumnList = updateColumnList;
      this.sourceExpressionList = sourceExpressionList;
      if (operation == Operation.UPDATE) {
        requireNonNull(updateColumnList, "updateColumnList");
        requireNonNull(sourceExpressionList, "sourceExpressionList");
        checkArgument(sourceExpressionList.size() == updateColumnList.size());
      } else {
        if (operation == Operation.MERGE) {
          requireNonNull(updateColumnList, "updateColumnList");
        } else {
          checkArgument(updateColumnList == null);
        }
        checkArgument(sourceExpressionList == null);
      }
      RelOptSchema relOptSchema = table.getRelOptSchema();
      if (relOptSchema != null) {
        cluster.getPlanner().registerSchema(relOptSchema);
      }
      this.flattened = flattened;
    }
  
    /**
     * Creates a TableModify by parsing serialized output.
     */
    protected TableModify(RelInput input) {
      this(input.getCluster(),
          input.getTraitSet(),
          input.getTable("table"),
          requireNonNull(
              (Prepare.CatalogReader) input.getTable("table").getRelOptSchema(),
              "relOptSchema"),
          input.getInput(),
          requireNonNull(input.getEnum("operation", Operation.class), "operation"),
          input.getStringList("updateColumnList"),
          input.getExpressionList("sourceExpressionList"),
          input.getBoolean("flattened", false));
    }
  
    //~ Methods ----------------------------------------------------------------
  
    public Prepare.CatalogReader getCatalogReader() {
      return catalogReader;
    }
  
    @Override public RelOptTable getTable() {
      return table;
    }
  
    public @Nullable List<String> getUpdateColumnList() {
      return updateColumnList;
    }
  
    public @Nullable List<RexNode> getSourceExpressionList() {
      return sourceExpressionList;
    }
  
    public boolean isFlattened() {
      return flattened;
    }
  
    public Operation getOperation() {
      return operation;
    }
  
    public boolean isInsert() {
      return operation == Operation.INSERT;
    }
  
    public boolean isUpdate() {
      return operation == Operation.UPDATE;
    }
  
    public boolean isDelete() {
      return operation == Operation.DELETE;
    }
  
    public boolean isMerge() {
      return operation == Operation.MERGE;
    }
  
    @Override public RelDataType deriveRowType() {
      return RelOptUtil.createDmlRowType(
          SqlKind.INSERT, getCluster().getTypeFactory());
    }
  
    @Override public RelDataType getExpectedInputRowType(int ordinalInParent) {
      assert ordinalInParent == 0;
  
      if (inputRowType != null) {
        return inputRowType;
      }
  
      final RelDataTypeFactory typeFactory = getCluster().getTypeFactory();
      final RelDataType rowType = table.getRowType();
      switch (operation) {
      case UPDATE:
        if (updateColumnList == null) {
          throw new AssertionError("updateColumnList must not be null for "
              + operation);
        }
        inputRowType =
            typeFactory.createJoinType(rowType,
                getCatalogReader().createTypeFromProjection(rowType,
                    updateColumnList));
        break;
      case MERGE:
        if (updateColumnList == null) {
          throw new AssertionError("updateColumnList must not be null for "
              + operation);
        }
        inputRowType =
            typeFactory.createJoinType(
                typeFactory.createJoinType(rowType, rowType),
                getCatalogReader().createTypeFromProjection(rowType,
                    updateColumnList));
        break;
      default:
        inputRowType = rowType;
        break;
      }
  
      if (flattened) {
        inputRowType =
            SqlTypeUtil.flattenRecordType(
                typeFactory,
                inputRowType,
                null);
      }
  
      return inputRowType;
    }
  
    @Override public RelWriter explainTerms(RelWriter pw) {
      return super.explainTerms(pw)
          .item("table", table.getQualifiedName())
          .item("operation", RelEnumTypes.fromEnum(getOperation()))
          .itemIf("updateColumnList", updateColumnList, updateColumnList != null)
          .itemIf("sourceExpressionList", sourceExpressionList,
              sourceExpressionList != null)
          .item("flattened", flattened);
    }
  
    @Override public @Nullable RelOptCost computeSelfCost(RelOptPlanner planner,
        RelMetadataQuery mq) {
      // REVIEW jvs 21-Apr-2006:  Just for now...
      double rowCount = mq.getRowCount(this);
      return planner.getCostFactory().makeCost(rowCount, 0, 0);
    }
  }
  