<commit_before>package org.labkey.dbutils.api;

import org.json.JSONObject;
import org.labkey.api.collections.CaseInsensitiveMapWrapper;
import org.labkey.api.data.Container;
import org.labkey.api.data.DbScope;
import org.labkey.api.data.TableInfo;
import org.labkey.api.query.BatchValidationException;
import org.labkey.api.query.DuplicateKeyException;
import org.labkey.api.query.InvalidKeyException;
import org.labkey.api.query.QueryService;
import org.labkey.api.query.QueryUpdateService;
import org.labkey.api.query.QueryUpdateServiceException;
import org.labkey.api.security.User;

import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * Similarly to how {@link SimpleQuery} makes it simple to access data in tables, {@link SimpleQueryUpdater} makes
 * it easy to edit tables.
 */
public class SimpleQueryUpdater {
    private TableInfo tableInfo;
    private QueryUpdateService service;
    private User user;
    private Container container;

    public SimpleQueryUpdater (User user, Container container, String schemaName, String tableName) {
        this.tableInfo = QueryService.get().getUserSchema(user, container, schemaName).getTable(tableName);
        this.service = this.tableInfo.getUpdateService();
        this.user = user;
        this.container = container;
    }

    public List<CaseInsensitiveMapWrapper<Object>> upsert(Map<String, Object>... rowArray) throws QueryUpdateServiceException, SQLException, InvalidKeyException, BatchValidationException, DuplicateKeyException {
        List<Map<String, Object>> rows = makeRowsCaseInsensitive(rowArray);
        return this.upsert(rows);
    }


    /**
     * Inserts or updates a row to ensure it's in the table.  It intelligently checks to see if each row exists,
     * and determines whether to insert or update them.  It does the entire thing as a transaction, so any failure
     * will result in no data being inserted.
     *
     * @param rows A "list" of {@link Map<String, Object>} objects that represent rows.  Note that these can
     *                 be {@link JSONObject}s, if you so choose.
     * @return A {@link List<CaseInsensitiveMapWrapper>} of the rows that were upserted.  They will have all been
     *         converted to be case-insensitive, and will include any fields that were added/changed automatically
     *         by the INSERT/UPDATE or trigger scripts.
     * @throws QueryUpdateServiceException Indicates an internal error in the QueryUpdateService
     * @throws SQLException Indicates a communication problem with the SQL database
     * @throws InvalidKeyException Should not occur, since any row that would throw this will be INSERTed instead
     * @throws BatchValidationException Thrown if the trigger scripts throw an error on any row.
     * @throws DuplicateKeyException Should not occur, since any row that would throw this will be UPDATEd instead
     */
    public List<CaseInsensitiveMapWrapper<Object>> upsert(List<Map<String, Object>> rows) throws QueryUpdateServiceException, SQLException, InvalidKeyException, BatchValidationException, DuplicateKeyException {
        rows = makeRowListCaseInsensitive(rows);
        List<CaseInsensitiveMapWrapper<Object>> rowsToReturn = new ArrayList<>();

        // Create buckets for our rows to insert or update
        List<Map<String, Object>> rowsToInsert = new ArrayList<>();
        List<Map<String, Object>> rowsToUpdate = new ArrayList<>();

        // Rows that can be found in the database
        List<Map<String, Object>> rowsInDB = service.getRows(user, container, rows);

        // Assign each row to insert/update queues
        rowLoop:
        for (Map<String, Object> row : rows) {
            for (Map<String, Object> rowFromDB : rowsInDB) {
                if (this.rowMapsAreEqual(rowFromDB, row)) {
                    rowsToUpdate.add(row);
                    continue rowLoop;
                }
            }

            rowsToInsert.add(row);
        }

        // Wrap everything in a transaction to make both UPDATE and INSERT operations atomic together.
        try(DbScope.Transaction transaction = tableInfo.getSchema().getScope().ensureTransaction()) {
            BatchValidationException validationException = new BatchValidationException();

            if (rowsToInsert.size() > 0) {
                List<Map<String, Object>> insertedRows = service.insertRows(user, container, rowsToInsert, validationException, null, null);

                if (validationException.hasErrors()) {
                    throw validationException;
                }
                // From my understanding, this shouldn't really happen, if you're checking the BatchValidationException.
                // Otherwise, the QueryUpdateService has a tendency to insert as much as it can, and not throw an error
                // for the rows that it couldn't.
                else if (insertedRows.size() != rowsToInsert.size()) {
                    throw new QueryUpdateServiceException("Not all rows were inserted properly");
                }
                else {
                    rowsToReturn.addAll(getCaseInsensitiveRowList(insertedRows));
                }
            }

            if (rowsToUpdate.size() > 0) {
                List<Map<String, Object>> updatedRows = service.updateRows(user, container, rowsToUpdate, rowsToUpdate, null, null);

                if (updatedRows.size() != rowsToUpdate.size()) {
                    throw new QueryUpdateServiceException("Not all rows updated properly");
                }
                else {
                    rowsToReturn.addAll(getCaseInsensitiveRowList(updatedRows));
                }
            }

            // Actually commit the changes.
            transaction.commit();
        }

        return rowsToReturn;
    }

    public List<CaseInsensitiveMapWrapper<Object>> delete(Map<String, Object>... rowArray) throws QueryUpdateServiceException, SQLException, InvalidKeyException, BatchValidationException, DuplicateKeyException {
        List<Map<String, Object>> rowsToDelete = makeRowsCaseInsensitive(rowArray);
        List<CaseInsensitiveMapWrapper<Object>> rowsToReturn = new ArrayList<>();

        // Wrap everything in a transaction to make both UPDATE and INSERT operations atomic together.
        try(DbScope.Transaction transaction = tableInfo.getSchema().getScope().ensureTransaction()) {
            if (rowsToDelete.size() > 0) {
                List<Map<String, Object>> deletedRows = service.deleteRows(user, container, rowsToDelete, null, null);

                // Check to make sure that the QueryUpdateService doesn't try to delete as much as it can, and not throw
                // an error for the rows that it couldn't.
                if (deletedRows.size() != deletedRows.size()) {
                    throw new QueryUpdateServiceException("Not all rows were deleted properly");
                }
                else {
                    rowsToReturn.addAll(getCaseInsensitiveRowList(deletedRows));
                }
            }
<commit_msg>added insert and update methods and refactored upsert method to remove duplicated code<commit_after>package org.labkey.dbutils.api;

import org.json.JSONObject;
import org.labkey.api.collections.CaseInsensitiveMapWrapper;
import org.labkey.api.data.Container;
import org.labkey.api.data.DbScope;
import org.labkey.api.data.TableInfo;
import org.labkey.api.query.BatchValidationException;
import org.labkey.api.query.DuplicateKeyException;
import org.labkey.api.query.InvalidKeyException;
import org.labkey.api.query.QueryService;
import org.labkey.api.query.QueryUpdateService;
import org.labkey.api.query.QueryUpdateServiceException;
import org.labkey.api.security.User;

import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * Similarly to how {@link SimpleQuery} makes it simple to access data in tables, {@link SimpleQueryUpdater} makes
 * it easy to edit tables.
 */
public class SimpleQueryUpdater {
    private TableInfo tableInfo;
    private QueryUpdateService service;
    private User user;
    private Container container;

    public SimpleQueryUpdater (User user, Container container, String schemaName, String tableName) {
        this.tableInfo = QueryService.get().getUserSchema(user, container, schemaName).getTable(tableName);
        this.service = this.tableInfo.getUpdateService();
        this.user = user;
        this.container = container;
    }

    public List<CaseInsensitiveMapWrapper<Object>> insert(Map<String, Object>... rowArray) throws QueryUpdateServiceException, SQLException, BatchValidationException, DuplicateKeyException {
        List<Map<String, Object>> rows = makeRowsCaseInsensitive(rowArray);
        return this.insert(rows);
    }

    public List<CaseInsensitiveMapWrapper<Object>> update(Map<String, Object>... rowArray) throws QueryUpdateServiceException, SQLException, BatchValidationException, InvalidKeyException {
        List<Map<String, Object>> rows = makeRowsCaseInsensitive(rowArray);
        return this.update(rows);
    }

    public List<CaseInsensitiveMapWrapper<Object>> upsert(Map<String, Object>... rowArray) throws QueryUpdateServiceException, SQLException, InvalidKeyException, BatchValidationException, DuplicateKeyException {
        List<Map<String, Object>> rows = makeRowsCaseInsensitive(rowArray);
        return this.upsert(rows);
    }

    /**
     * Inserts a row into the table. It does the entire thing as a transaction, so any failure
     * will result in no data being inserted.
     *
     * @param rows A "list" of {@link Map<String, Object>} objects that represent rows.  Note that these can
     *      be {@link JSONObject}s, if you so choose.
     * @return A {@link List<CaseInsensitiveMapWrapper>} of the rows that were inserted.  They will have all been
     *      converted to be case-insensitive, and will include any fields that were added/changed automatically
     *      by the INSERT or trigger scripts.
     * @throws QueryUpdateServiceException Indicates an internal error in the QueryUpdateService
     * @throws SQLException Indicates a communication problem with the SQL database
     * @throws BatchValidationException Thrown if the trigger scripts throw an error on any row
     * @throws DuplicateKeyException Thrown if there is already a row in the table with the same key
     */
    public List<CaseInsensitiveMapWrapper<Object>> insert(List<Map<String, Object>> rows) throws QueryUpdateServiceException, SQLException, BatchValidationException, DuplicateKeyException {
        rows = makeRowListCaseInsensitive(rows);
        List<CaseInsensitiveMapWrapper<Object>> rowsToReturn = new ArrayList<>();

        // Wrap everything in a transaction to make INSERT operations atomic together.
        try(DbScope.Transaction transaction = tableInfo.getSchema().getScope().ensureTransaction()) {
            rowsToReturn.addAll(getCaseInsensitiveRowList(doInsert(rows)));

            // Actually commit the changes.
            transaction.commit();
        }
        return rowsToReturn;
    }

    /**
     * Updates a row in the table. It does the entire thing as a transaction, so any failure
     * will result in no data being updated.
     *
     * @param rows A "list" of {@link Map<String, Object>} objects that represent rows.  Note that these can
     *      be {@link JSONObject}s, if you so choose.
     * @return A {@link List<CaseInsensitiveMapWrapper>} of the rows that were updated.  They will have all been
     *      converted to be case-insensitive, and will include any fields that were added/changed automatically
     *      by the UPDATE or trigger scripts.
     * @throws QueryUpdateServiceException Indicates an internal error in the QueryUpdateService
     * @throws SQLException Indicates a communication problem with the SQL database
     * @throws InvalidKeyException Thrown if the key in an updated row doesn't exist in the table
     * @throws BatchValidationException Thrown if the trigger scripts throw an error on any row
     */
    public List<CaseInsensitiveMapWrapper<Object>> update(List<Map<String, Object>> rows) throws QueryUpdateServiceException, SQLException, BatchValidationException, InvalidKeyException {
        rows = makeRowListCaseInsensitive(rows);
        List<CaseInsensitiveMapWrapper<Object>> rowsToReturn = new ArrayList<>();

        // Wrap everything in a transaction to make UPDATE operations atomic together.
        try(DbScope.Transaction transaction = tableInfo.getSchema().getScope().ensureTransaction()) {
            rowsToReturn.addAll(getCaseInsensitiveRowList(doUpdate(rows)));

            // Actually commit the changes.
            transaction.commit();
        }
        return rowsToReturn;
    }



    /**
     * Inserts or updates a row to ensure it's in the table.  It intelligently checks to see if each row exists,
     * and determines whether to insert or update them.  It does the entire thing as a transaction, so any failure
     * will result in no data being inserted.
     *
     * @param rows A "list" of {@link Map<String, Object>} objects that represent rows.  Note that these can
     *                 be {@link JSONObject}s, if you so choose.
     * @return A {@link List<CaseInsensitiveMapWrapper>} of the rows that were upserted.  They will have all been
     *         converted to be case-insensitive, and will include any fields that were added/changed automatically
     *         by the INSERT/UPDATE or trigger scripts.
     * @throws QueryUpdateServiceException Indicates an internal error in the QueryUpdateService
     * @throws SQLException Indicates a communication problem with the SQL database
     * @throws InvalidKeyException Should not occur, since any row that would throw this will be INSERTed instead
     * @throws BatchValidationException Thrown if the trigger scripts throw an error on any row.
     * @throws DuplicateKeyException Should not occur, since any row that would throw this will be UPDATEd instead
     */
    public List<CaseInsensitiveMapWrapper<Object>> upsert(List<Map<String, Object>> rows) throws QueryUpdateServiceException, SQLException, InvalidKeyException, BatchValidationException, DuplicateKeyException {
        rows = makeRowListCaseInsensitive(rows);
        List<CaseInsensitiveMapWrapper<Object>> rowsToReturn = new ArrayList<>();

        // Create buckets for our rows to insert or update
        List<Map<String, Object>> rowsToInsert = new ArrayList<>();
        List<Map<String, Object>> rowsToUpdate = new ArrayList<>();

        // Rows that can be found in the database
        List<Map<String, Object>> rowsInDB = service.getRows(user, container, rows);

        // Assign each row to insert/update queues
        rowLoop:
        for (Map<String, Object> row : rows) {
            for (Map<String, Object> rowFromDB : rowsInDB) {
                if (this.rowMapsAreEqual(rowFromDB, row)) {
                    rowsToUpdate.add(row);
                    continue rowLoop;
                }
            }

            rowsToInsert.add(row);
        }

        // Wrap everything in a transaction to make both UPDATE and INSERT operations atomic together.
        try(DbScope.Transaction transaction = tableInfo.getSchema().getScope().ensureTransaction()) {
            rowsToReturn.addAll(getCaseInsensitiveRowList(doInsert(rowsToInsert)));
            rowsToReturn.addAll(getCaseInsensitiveRowList(doUpdate(rowsToUpdate)));

            // Actually commit the changes.
            transaction.commit();
        }
        return rowsToReturn;
    }

    public List<Map<String, Object>> doInsert(List<Map<String, Object>> rows) throws BatchValidationException , DuplicateKeyException, SQLException, QueryUpdateServiceException{
        List<Map<String, Object>> insertedRows = null;
        if (rows.size() > 0) {
            BatchValidationException validationException = new BatchValidationException();
            insertedRows = service.insertRows(user, container, rows, validationException, null, null);

            if (validationException.hasErrors()) {
                throw validationException;
            }
            // From my understanding, this shouldn't really happen, if you're checking the BatchValidationException.
            // Otherwise, the QueryUpdateService has a tendency to insert as much as it can, and not throw an error
            // for the rows that it couldn't.
            else if (insertedRows.size() != rows.size()) {
                throw new QueryUpdateServiceException("Not all rows were inserted properly");
            }
        }
        return insertedRows;
    }

    public List<Map<String, Object>> doUpdate(List<Map<String, Object>> rows) throws BatchValidationException , InvalidKeyException, SQLException, QueryUpdateServiceException {
        List<Map<String, Object>> updatedRows = null;
        if (rows.size() > 0) {
            updatedRows = service.updateRows(user, container, rows, rows, null, null);

            if (updatedRows.size() != rows.size()) {
                throw new QueryUpdateServiceException("Not all rows updated properly");
            }
        }
        return updatedRows;
    }

    public List<CaseInsensitiveMapWrapper<Object>> delete(Map<String, Object>... rowArray) throws QueryUpdateServiceException, SQLException, InvalidKeyException, BatchValidationException {
        List<Map<String, Object>> rowsToDelete = makeRowsCaseInsensitive(rowArray);
        List<CaseInsensitiveMapWrapper<Object>> rowsToReturn = new ArrayList<>();

        // Wrap everything in a transaction to make both UPDATE and INSERT operations atomic together.
        try(DbScope.Transaction transaction = tableInfo.getSchema().getScope().ensureTransaction()) {
            if (rowsToDelete.size() > 0) {
                List<Map<String, Object>> deletedRows = service.deleteRows(user, container, rowsToDelete, null, null);

                // Check to make sure that the QueryUpdateService doesn't try to delete as much as it can, and not throw
                // an error for the rows that it couldn't.
                if (deletedRows.size() != deletedRows.size()) {
                    throw new QueryUpdateServiceException("Not all rows were deleted properly");
                }
                else {
                    rowsToReturn.addAll(getCaseInsensitiveRowList(deletedRows));
                }
            }

{'fraction_non_alphanumeric': 0.08110002696144514, 'fraction_numerical': 0.0003235373416015098, 'mean_word_length': 3.0334928229665072, 'pattern_counts': {'":': 0, '<': 115, '<?xml version=': 0, '>': 121, 'https://': 0, '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': '2767741', 'n_tokens_mistral': 4631, 'n_tokens_neox': 4500, 'n_words': 1514}