<commit_before>package sqlancer.tidb.gen;

import java.sql.SQLException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import sqlancer.Query;
import sqlancer.QueryAdapter;
import sqlancer.Randomly;
import sqlancer.tidb.TiDBExpressionGenerator;
import sqlancer.tidb.TiDBProvider.TiDBGlobalState;
import sqlancer.tidb.TiDBSchema.TiDBColumn;
import sqlancer.tidb.TiDBSchema.TiDBTable;

public class TiDBInsertGenerator {

	private final TiDBGlobalState globalState;
	private final Set<String> errors = new HashSet<>();
	private TiDBExpressionGenerator gen;

	public TiDBInsertGenerator(TiDBGlobalState globalState) {
		this.globalState = globalState;
		errors.add("Duplicate entry");
		errors.add("cannot be null");
		errors.add("doesn't have a default value");
		errors.add("Out of range value");
		errors.add("Incorrect tinyint value");
		errors.add("Data truncation");
		errors.add("Bad Number");
		errors.add("The value specified for generated column"); // TODO: do not insert data into generated columns
	}

	public static Query getQuery(TiDBGlobalState globalState) throws SQLException {
		return new TiDBInsertGenerator(globalState).get();
	}

	private Query get() {
		TiDBTable table = globalState.getSchema().getRandomTable(t -> !t.isView());
		gen = new TiDBExpressionGenerator(globalState).setColumns(table.getColumns());
		StringBuilder sb = new StringBuilder();
		sb.append(Randomly.fromOptions("INSERT", "REPLACE"));
		sb.append(" INTO ");
		sb.append(table.getName());
		if (Randomly.getBoolean()) {
			sb.append(" VALUES ");
			List<TiDBColumn> columns = table.getColumns();
			insertColumns(sb, columns);
		} else {
			List<TiDBColumn> columnSubset = table.getRandomNonEmptyColumnSubset();
			sb.append("(");
			sb.append(columnSubset.stream().map(c -> c.getName()).collect(Collectors.joining(", ")));
			sb.append(") VALUES ");
			insertColumns(sb, columnSubset);
		}
		return new QueryAdapter(sb.toString(), errors);
	}

	private void insertColumns(StringBuilder sb, List<TiDBColumn> columns) {
		for (int nrRows = 0; nrRows < Randomly.smallNumber() + 1; nrRows++) {
			if (nrRows != 0) {
				sb.append(", ");
			}
			sb.append("(");
			for (int nrColumn = 0; nrColumn < columns.size(); nrColumn++) {
				if (nrColumn != 0) {
					sb.append(", ");
				}
				insertValue(sb, columns.get(nrColumn));
			}
			sb.append(")");
		}
	}
<commit_msg>[TiDB] Enhance the INSERT statement<commit_after>package sqlancer.tidb.gen;

import java.sql.SQLException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import sqlancer.Query;
import sqlancer.QueryAdapter;
import sqlancer.Randomly;
import sqlancer.tidb.TiDBExpressionGenerator;
import sqlancer.tidb.TiDBProvider.TiDBGlobalState;
import sqlancer.tidb.TiDBSchema.TiDBColumn;
import sqlancer.tidb.TiDBSchema.TiDBTable;
import sqlancer.tidb.visitor.TiDBVisitor;

public class TiDBInsertGenerator {

	private final TiDBGlobalState globalState;
	private final Set<String> errors = new HashSet<>();
	private TiDBExpressionGenerator gen;

	public TiDBInsertGenerator(TiDBGlobalState globalState) {
		this.globalState = globalState;
		errors.add("Duplicate entry");
		errors.add("cannot be null");
		errors.add("doesn't have a default value");
		errors.add("Out of range value");
		errors.add("Incorrect tinyint value");
		errors.add("Data truncation");
		errors.add("Bad Number");
		errors.add("The value specified for generated column"); // TODO: do not insert data into generated columns
	}

	public static Query getQuery(TiDBGlobalState globalState) throws SQLException {
		return new TiDBInsertGenerator(globalState).get();
	}

	private Query get() {
		TiDBTable table = globalState.getSchema().getRandomTable(t -> !t.isView());
		gen = new TiDBExpressionGenerator(globalState).setColumns(table.getColumns());
		StringBuilder sb = new StringBuilder();
		boolean isInsert = Randomly.getBoolean();
		if (isInsert) {
			sb.append("INSERT");
		} else {
			sb.append("REPLACE");
		}
		if (Randomly.getBooleanWithRatherLowProbability()) {
			sb.append(" ");
			sb.append(Randomly.fromOptions("LOW_PRIORITY", "HIGH_PRIORITY", "DELAYED"));
		}
		if (isInsert && Randomly.getBoolean() && false /* https://github.com/pingcap/tidb/issues/16025 */) {
			sb.append(" IGNORE ");
		}
		sb.append(" INTO ");
		sb.append(table.getName());
		if (Randomly.getBoolean()) {
			sb.append(" VALUES ");
			List<TiDBColumn> columns = table.getColumns();
			insertColumns(sb, columns);
		} else {
			List<TiDBColumn> columnSubset = table.getRandomNonEmptyColumnSubset();
			sb.append("(");
			sb.append(columnSubset.stream().map(c -> c.getName()).collect(Collectors.joining(", ")));
			sb.append(") VALUES ");
			insertColumns(sb, columnSubset);
		}
		if (isInsert && Randomly.getBoolean()) {
			sb.append(" ON DUPLICATE KEY UPDATE ");
			sb.append(table.getRandomColumn().getName());
			sb.append("=");
			sb.append(TiDBVisitor.asString(gen.generateExpression()));
		}
		return new QueryAdapter(sb.toString(), errors);
	}

	private void insertColumns(StringBuilder sb, List<TiDBColumn> columns) {
		for (int nrRows = 0; nrRows < Randomly.smallNumber() + 1; nrRows++) {
			if (nrRows != 0) {
				sb.append(", ");
			}
			sb.append("(");
			for (int nrColumn = 0; nrColumn < columns.size(); nrColumn++) {
				if (nrColumn != 0) {
					sb.append(", ");
				}
				insertValue(sb, columns.get(nrColumn));
			}
			sb.append(")");
		}
	}

{'fraction_non_alphanumeric': 0.14627028998723327, 'fraction_numerical': 0.0027357286157213205, 'mean_word_length': 4.790918690601901, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5208842', 'n_tokens_mistral': 1951, 'n_tokens_neox': 1733, 'n_words': 372}