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.cassandra.harry.execution;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import accord.utils.Invariants;
  import org.apache.cassandra.harry.SchemaSpec;
  import org.apache.cassandra.harry.cql.DeleteHelper;
  import org.apache.cassandra.harry.cql.SelectHelper;
  import org.apache.cassandra.harry.cql.WriteHelper;
  import org.apache.cassandra.harry.op.Operations;
  import org.apache.cassandra.harry.op.Visit;
  
  // TODO: this class can be substantially improved by removing internal mutable state
  public class QueryBuildingVisitExecutor extends VisitExecutor
  {
      private static final Logger logger = LoggerFactory.getLogger(QueryBuildingVisitExecutor.class);
      protected final SchemaSpec schema;
      protected final WrapQueries wrapQueries;
  
      public QueryBuildingVisitExecutor(SchemaSpec schema, WrapQueries wrapQueries)
      {
          this.schema = schema;
          this.wrapQueries = wrapQueries;
      }
  
      public BuiltQuery compile(Visit visit)
      {
          beginLts(visit.lts);
          for (Operations.Operation op : visit.operations)
          {
              if (logger.isTraceEnabled())
                  logger.trace("{} {}", visit.lts, op);
              operation(op);
          }
  
          // TODO: try inducing timeouts and checking non-propagation or discovery
          endLts(visit.lts);
  
          if (!statements.isEmpty())
          {
              Object[] bindingsArray = new Object[bindings.size()];
              bindings.toArray(bindingsArray);
              BuiltQuery query = new BuiltQuery(selects,
                                     wrapQueries.wrap(visit, String.join("\n ", statements)),
                                     bindingsArray);
              clear();
              return query;
          }
  
          Invariants.require(bindings.isEmpty() && visitedPds.isEmpty() && selects.isEmpty());
          return null;
      }
  
      public static class BuiltQuery extends CompiledStatement
      {
          protected final List<Operations.SelectStatement> selects;
          public BuiltQuery(List<Operations.SelectStatement> selects, String cql, Object... bindings)
          {
              super(cql, bindings);
              this.selects = selects;
          }
      }
  
      /**
       * Per-LTS state
       */
      private final List<String> statements = new ArrayList<>();
      private final List<Object> bindings = new ArrayList<>();
      private final Set<Long> visitedPds = new HashSet<>();
  
      private List<Operations.SelectStatement> selects = null;
  
      protected void beginLts(long lts)
      {
          statements.clear();
          bindings.clear();
          visitedPds.clear();
          selects = new ArrayList<>();
      }
  
      protected void endLts(long lts)
      {
          if (statements.isEmpty())
          {
              Invariants.require(bindings.isEmpty() && visitedPds.isEmpty() && selects.isEmpty());
              return;
          }
  
          assert visitedPds.size() == 1 : String.format("Token aware only works with a single value per token, but got %s", visitedPds);
      }
  
      private void clear()
      {
          statements.clear();
          bindings.clear();
          visitedPds.clear();
          selects = null;
      }
  
      protected void operation(Operations.Operation operation)
      {
          if (operation instanceof Operations.PartitionOperation)
              visitedPds.add(((Operations.PartitionOperation) operation).pd());
          CompiledStatement statement;
          switch (operation.kind())
          {
              case UPDATE:
                  statement = WriteHelper.inflateUpdate((Operations.WriteOp) operation, schema, operation.lts());
                  break;
              case INSERT:
                  statement = WriteHelper.inflateInsert((Operations.WriteOp) operation, schema, operation.lts());
                  break;
              case DELETE_RANGE:
                  statement = DeleteHelper.inflateDelete((Operations.DeleteRange) operation, schema, operation.lts());
                  break;
              case DELETE_PARTITION:
                  statement = DeleteHelper.inflateDelete((Operations.DeletePartition) operation, schema, operation.lts());
                  break;
              case DELETE_ROW:
                  statement = DeleteHelper.inflateDelete((Operations.DeleteRow) operation, schema, operation.lts());
                  break;
              case DELETE_COLUMNS:
                  statement = DeleteHelper.inflateDelete((Operations.DeleteColumns) operation, schema, operation.lts());
                  break;
              case SELECT_PARTITION:
                  statement = SelectHelper.select((Operations.SelectPartition) operation, schema);
                  selects.add((Operations.SelectStatement) operation);
                  break;
              case SELECT_ROW:
                  statement = SelectHelper.select((Operations.SelectRow) operation, schema);
                  selects.add((Operations.SelectStatement) operation);
                  break;
              case SELECT_RANGE:
                  statement = SelectHelper.select((Operations.SelectRange) operation, schema);
                  selects.add((Operations.SelectStatement) operation);
                  break;
              case SELECT_CUSTOM:
                  statement = SelectHelper.select((Operations.SelectCustom) operation, schema);
                  selects.add((Operations.SelectStatement) operation);
                  break;
  
              case CUSTOM:
                  ((Operations.CustomRunnableOperation) operation).execute();
                  return;
              default:
                  throw new IllegalArgumentException();
          }
          statements.add(statement.cql());
          Collections.addAll(bindings, statement.bindings());
      }
  
      private static final String wrapInUnloggedBatchFormat = "BEGIN UNLOGGED BATCH\n" +
                                                              "  %s\n" +
                                                              "APPLY BATCH;";
  
      private static final String wrapInTxnFormat = "BEGIN TRANSACTION\n" +
                                                    "  %s\n" +
                                                    "COMMIT TRANSACTION;";
  
      public interface WrapQueries
      {
          WrapQueries UNLOGGED_BATCH = (visit, compiled) -> {
              if (visit.operations.length == 1)
                  return compiled;
              return String.format(wrapInUnloggedBatchFormat, compiled);
          };
  
          WrapQueries TRANSACTION = (visit, compiled) -> String.format(wrapInTxnFormat, compiled);
  
  
          String wrap(Visit visit, String compiled);
      }
  }
  