<commit_before>/*
 * Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
 * license agreements.  See the NOTICE file distributed with this work for
 * additional information regarding copyright ownership.  Crate 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.
 *
 * However, if you have executed another commercial license agreement
 * with Crate these terms will supersede the license and you may use the
 * software solely pursuant to the terms of the relevant commercial agreement.
 */

package io.crate.action.sql;

import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import io.crate.Constants;
import io.crate.analyze.*;
import io.crate.blob.v2.BlobIndices;
import io.crate.core.collections.Bucket;
import io.crate.core.collections.Row;
import io.crate.exceptions.AlterTableAliasException;
import io.crate.executor.Executor;
import io.crate.executor.Job;
import io.crate.executor.TaskResult;
import io.crate.executor.transport.TransportActionProvider;
import io.crate.metadata.PartitionName;
import io.crate.metadata.table.TableInfo;
import io.crate.operation.aggregation.impl.CountAggregation;
import io.crate.planner.IterablePlan;
import io.crate.planner.Plan;
import io.crate.planner.RowGranularity;
import io.crate.planner.node.dql.CollectNode;
import io.crate.planner.node.dql.MergeNode;
import io.crate.planner.projection.AggregationProjection;
import io.crate.planner.projection.Projection;
import io.crate.planner.symbol.Aggregation;
import io.crate.planner.symbol.InputColumn;
import io.crate.planner.symbol.Symbol;
import io.crate.types.DataType;
import io.crate.types.DataTypes;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.admin.indices.alias.Alias;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingResponse;
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
import org.elasticsearch.action.admin.indices.refresh.RefreshResponse;
import org.elasticsearch.action.admin.indices.settings.put.UpdateSettingsRequest;
import org.elasticsearch.action.admin.indices.settings.put.UpdateSettingsResponse;
import org.elasticsearch.action.admin.indices.template.get.GetIndexTemplatesRequest;
import org.elasticsearch.action.admin.indices.template.get.GetIndexTemplatesResponse;
import org.elasticsearch.action.admin.indices.template.put.PutIndexTemplateRequest;
import org.elasticsearch.action.admin.indices.template.put.PutIndexTemplateResponse;
import org.elasticsearch.cluster.ClusterService;
import org.elasticsearch.cluster.metadata.AliasMetaData;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.cluster.metadata.IndexTemplateMetaData;
import org.elasticsearch.common.compress.CompressedString;
import org.elasticsearch.common.inject.Inject;
import org.elasticsearch.common.inject.Provider;
import org.elasticsearch.common.settings.ImmutableSettings;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentHelper;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * visitor that dispatches requests based on Analysis class to different actions.
 *
 * Its methods return a future returning a Long containing the response rowCount.
 * If the future returns <code>null</code>, no row count shall be created.
 */
public class DDLStatementDispatcher extends AnalyzedStatementVisitor<Void, ListenableFuture<Long>> {

    private final ClusterService clusterService;
    private final BlobIndices blobIndices;
    private final Provider<Executor> executorProvider;
    private final TransportActionProvider transportActionProvider;

    @Inject
    public DDLStatementDispatcher(ClusterService clusterService,
                                  BlobIndices blobIndices,
                                  Provider<Executor> executorProvider,
                                  TransportActionProvider transportActionProvider) {
        this.clusterService = clusterService;
        this.blobIndices = blobIndices;
        this.executorProvider = executorProvider;
        this.transportActionProvider = transportActionProvider;
    }

    @Override
    protected ListenableFuture<Long> visitAnalyzedStatement(AnalyzedStatement analyzedStatement, Void context) {
        throw new UnsupportedOperationException(String.format("Can't handle \"%s\"", analyzedStatement));
    }

   @Override
    public ListenableFuture<Long> visitCreateBlobTableStatement(
           CreateBlobTableAnalyzedStatement analysis, Void context) {
        return wrapRowCountFuture(
                blobIndices.createBlobTable(
                        analysis.tableName(),
                        analysis.tableParameter().settings()
                ),
                1L
        );
    }

    @Override
    public ListenableFuture<Long> visitAddColumnStatement(final AddColumnAnalyzedStatement analysis, Void context) {
        final SettableFuture<Long> result = SettableFuture.create();
        if (analysis.newPrimaryKeys()) {
            Plan plan = genCountStarPlan(analysis.table());
            Job job = executorProvider.get().newJob(plan);
            ListenableFuture<List<TaskResult>> resultFuture = Futures.allAsList(executorProvider.get().execute(job));
            Futures.addCallback(resultFuture, new FutureCallback<List<TaskResult>>() {
                @Override
                public void onSuccess(@Nullable List<TaskResult> resultList) {
                    assert resultList != null && resultList.size() == 1;
                    Bucket rows = resultList.get(0).rows();
                    Row row = Iterables.getOnlyElement(rows);
                    if ((Long) row.get(0) == 0L) {
                        addColumnToTable(analysis, result);
                    } else {
                        result.setException(new UnsupportedOperationException(
                                "Cannot add a primary key column to a table that isn't empty"));
                    }
                }

                @Override
                public void onFailure(@Nonnull Throwable t) {
                    result.setException(t);
                }
            });
        } else {
            addColumnToTable(analysis, result);
        }
        return result;
    }

    private Plan genCountStarPlan(TableInfo table) {
        Aggregation countAggregationPartial = new Aggregation(
                CountAggregation.COUNT_STAR_FUNCTION,
                ImmutableList.<Symbol>of(),
                Aggregation.Step.ITER,
                Aggregation.Step.PARTIAL);
        Aggregation countAggregationFinal = new Aggregation(
                CountAggregation.COUNT_STAR_FUNCTION,
                ImmutableList.<Symbol>of(new InputColumn(0, DataTypes.LONG)),
                Aggregation.Step.PARTIAL,
                Aggregation.Step.FINAL);

        CollectNode collectNode = new CollectNode(
                "count",
                table.getRouting(WhereClause.MATCH_ALL, null),
                ImmutableList.<Symbol>of(),
                Arrays.<Projection>asList(new AggregationProjection(ImmutableList.of(countAggregationPartial))));
        collectNode.maxRowGranularity(RowGranularity.DOC);
        collectNode.outputTypes(ImmutableList.<DataType>of(DataTypes.UNDEFINED));
        MergeNode mergeNode = new MergeNode("local count merge", collectNode.executionNodes().size());
        mergeNode.projections(ImmutableList.<Projection>of(new AggregationProjection(ImmutableList.of(countAggregationFinal))));
        return new IterablePlan(collectNode, mergeNode);
    }

    private void addColumnToTable(AddColumnAnalyzedStatement analysis, final SettableFuture<Long> result) {
        boolean updateTemplate = analysis.table().isPartitioned();
        final AtomicInteger operations = new AtomicInteger(updateTemplate ? 2 : 1);
        final Map<String, Object> mapping = analysis.analyzedTableElements().toMapping();

        if (updateTemplate) {
            String templateName = PartitionName.templateName(analysis.table().ident().schema(), analysis.table().ident().name());
            IndexTemplateMetaData indexTemplateMetaData =
                    clusterService.state().metaData().templates().get(templateName);
            if (indexTemplateMetaData == null) {
                result.setException(new RuntimeException("Template for partitioned table is missing"));
            }
            mergeMappingAndUpdateTemplate(result, mapping, indexTemplateMetaData, operations);
        }

        // need to merge the _meta part of the mapping mapping before-hand because ES doesn't
        // update the _meta column recursively. Instead it is overwritten and therefore partitioned by
        // and collection_type information would be lost.
        String[] indexNames = getIndexNames(analysis.table(), null);
        if (indexNames.length == 0) {
            // if there are no indices yet we can return because we don't need to update existing mapping
            if (operations.decrementAndGet() == 0) {
                result.set(1L);
            }
            return;
        }
        PutMappingRequest request = new PutMappingRequest();
        request.indices(indexNames);
        request.type(Constants.DEFAULT_MAPPING_TYPE);<commit_msg>use planner for creating a count(*) plan at DDL dispatcher<commit_after>/*
 * Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
 * license agreements.  See the NOTICE file distributed with this work for
 * additional information regarding copyright ownership.  Crate 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.
 *
 * However, if you have executed another commercial license agreement
 * with Crate these terms will supersede the license and you may use the
 * software solely pursuant to the terms of the relevant commercial agreement.
 */

package io.crate.action.sql;

import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import io.crate.Constants;
import io.crate.analyze.*;
import io.crate.analyze.relations.TableRelation;
import io.crate.blob.v2.BlobIndices;
import io.crate.core.collections.Bucket;
import io.crate.core.collections.Row;
import io.crate.exceptions.AlterTableAliasException;
import io.crate.executor.Executor;
import io.crate.executor.Job;
import io.crate.executor.TaskResult;
import io.crate.executor.transport.TransportActionProvider;
import io.crate.metadata.OutputName;
import io.crate.metadata.PartitionName;
import io.crate.metadata.table.TableInfo;
import io.crate.operation.aggregation.impl.CountAggregation;
import io.crate.planner.Plan;
import io.crate.planner.Planner;
import io.crate.planner.symbol.Function;
import io.crate.planner.symbol.Symbol;
import io.crate.sql.tree.QualifiedName;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.admin.indices.alias.Alias;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingResponse;
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
import org.elasticsearch.action.admin.indices.refresh.RefreshResponse;
import org.elasticsearch.action.admin.indices.settings.put.UpdateSettingsRequest;
import org.elasticsearch.action.admin.indices.settings.put.UpdateSettingsResponse;
import org.elasticsearch.action.admin.indices.template.get.GetIndexTemplatesRequest;
import org.elasticsearch.action.admin.indices.template.get.GetIndexTemplatesResponse;
import org.elasticsearch.action.admin.indices.template.put.PutIndexTemplateRequest;
import org.elasticsearch.action.admin.indices.template.put.PutIndexTemplateResponse;
import org.elasticsearch.cluster.ClusterService;
import org.elasticsearch.cluster.metadata.AliasMetaData;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.cluster.metadata.IndexTemplateMetaData;
import org.elasticsearch.common.compress.CompressedString;
import org.elasticsearch.common.inject.Inject;
import org.elasticsearch.common.inject.Provider;
import org.elasticsearch.common.settings.ImmutableSettings;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentHelper;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * visitor that dispatches requests based on Analysis class to different actions.
 *
 * Its methods return a future returning a Long containing the response rowCount.
 * If the future returns <code>null</code>, no row count shall be created.
 */
public class DDLStatementDispatcher extends AnalyzedStatementVisitor<Void, ListenableFuture<Long>> {

    private final ClusterService clusterService;
    private final BlobIndices blobIndices;
    private final Provider<Executor> executorProvider;
    private final TransportActionProvider transportActionProvider;
    private final Planner planner;

    @Inject
    public DDLStatementDispatcher(ClusterService clusterService,
                                  BlobIndices blobIndices,
                                  Provider<Executor> executorProvider,
                                  TransportActionProvider transportActionProvider,
                                  Planner planner) {
        this.clusterService = clusterService;
        this.blobIndices = blobIndices;
        this.executorProvider = executorProvider;
        this.transportActionProvider = transportActionProvider;
        this.planner = planner;
    }

    @Override
    protected ListenableFuture<Long> visitAnalyzedStatement(AnalyzedStatement analyzedStatement, Void context) {
        throw new UnsupportedOperationException(String.format("Can't handle \"%s\"", analyzedStatement));
    }

   @Override
    public ListenableFuture<Long> visitCreateBlobTableStatement(
           CreateBlobTableAnalyzedStatement analysis, Void context) {
        return wrapRowCountFuture(
                blobIndices.createBlobTable(
                        analysis.tableName(),
                        analysis.tableParameter().settings()
                ),
                1L
        );
    }

    @Override
    public ListenableFuture<Long> visitAddColumnStatement(final AddColumnAnalyzedStatement analysis, Void context) {
        final SettableFuture<Long> result = SettableFuture.create();
        if (analysis.newPrimaryKeys()) {
            Plan plan = genCountStarPlan(analysis.table());
            Job job = executorProvider.get().newJob(plan);
            ListenableFuture<List<TaskResult>> resultFuture = Futures.allAsList(executorProvider.get().execute(job));
            Futures.addCallback(resultFuture, new FutureCallback<List<TaskResult>>() {
                @Override
                public void onSuccess(@Nullable List<TaskResult> resultList) {
                    assert resultList != null && resultList.size() == 1;
                    Bucket rows = resultList.get(0).rows();
                    Row row = Iterables.getOnlyElement(rows);
                    if ((Long) row.get(0) == 0L) {
                        addColumnToTable(analysis, result);
                    } else {
                        result.setException(new UnsupportedOperationException(
                                "Cannot add a primary key column to a table that isn't empty"));
                    }
                }

                @Override
                public void onFailure(@Nonnull Throwable t) {
                    result.setException(t);
                }
            });
        } else {
            addColumnToTable(analysis, result);
        }
        return result;
    }

    private Plan genCountStarPlan(TableInfo table) {
        QuerySpec querySpec = new QuerySpec();
        querySpec.where(WhereClause.MATCH_ALL);
        Function countFunction = new Function(
                CountAggregation.COUNT_STAR_FUNCTION, ImmutableList.<Symbol>of());
        querySpec.outputs(ImmutableList.<Symbol>of(countFunction));
        querySpec.hasAggregates(true);

        QueriedTable queriedTable = new QueriedTable(
                QualifiedName.of(table.ident().fqn()),
                new TableRelation(table),
                ImmutableList.of(new OutputName("count(*)")),
                querySpec
                );
        SelectAnalyzedStatement statement = new SelectAnalyzedStatement(queriedTable);
        return planner.process(statement, new Planner.Context());
    }

    private void addColumnToTable(AddColumnAnalyzedStatement analysis, final SettableFuture<Long> result) {
        boolean updateTemplate = analysis.table().isPartitioned();
        final AtomicInteger operations = new AtomicInteger(updateTemplate ? 2 : 1);
        final Map<String, Object> mapping = analysis.analyzedTableElements().toMapping();

        if (updateTemplate) {
            String templateName = PartitionName.templateName(analysis.table().ident().schema(), analysis.table().ident().name());
            IndexTemplateMetaData indexTemplateMetaData =
                    clusterService.state().metaData().templates().get(templateName);
            if (indexTemplateMetaData == null) {
                result.setException(new RuntimeException("Template for partitioned table is missing"));
            }
            mergeMappingAndUpdateTemplate(result, mapping, indexTemplateMetaData, operations);
        }

        // need to merge the _meta part of the mapping mapping before-hand because ES doesn't
        // update the _meta column recursively. Instead it is overwritten and therefore partitioned by
        // and collection_type information would be lost.
        String[] indexNames = getIndexNames(analysis.table(), null);
        if (indexNames.length == 0) {
            // if there are no indices yet we can return because we don't need to update existing mapping
            if (operations.decrementAndGet() == 0) {
                result.set(1L);
            }
            return;
        }
        PutMappingRequest request = new PutMappingRequest();
        request.indices(indexNames);
        request.type(Constants.DEFAULT_MAPPING_TYPE);
{'fraction_non_alphanumeric': 0.08506503214232322, 'fraction_numerical': 0.0015448248367967308, 'mean_word_length': 3.244500846023689, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2653081', 'n_tokens_mistral': 4978, 'n_tokens_neox': 4935, 'n_words': 1204}