SQLite.BulkOperations.cs

using System;
using System.Collections.Generic;
using System.Data;
using System.Threading;
using System.Threading.Tasks;
using DBAClientX.DataMovement;
using Microsoft.Data.Sqlite;

namespace DBAClientX;

public partial class SQLite
{
    private const int DefaultBulkInsertBatchSize = 500;
    internal const int BulkInsertParameterLimit = 999;

    private void BulkInsertCore(
        string database,
        DataTable table,
        string destinationTable,
        bool useTransaction,
        int? batchSize,
        bool useConnectionString)
    {
        ValidateBulkInsertInputs(table, destinationTable, batchSize);

        var connectionString = ResolveBulkInsertConnectionString(database, useConnectionString);

        SqliteConnection? connection = null;
        var dispose = false;
        SqliteTransaction? activeTransaction = null;
        SqliteTransaction? transaction = null;
        try
        {
            (connection, activeTransaction, dispose) = ResolveConnection(connectionString, useTransaction);

            if (!useTransaction)
            {
                transaction = connection.BeginTransaction();
            }

            var totalRows = table.Rows.Count;
            if (totalRows == 0)
            {
                return;
            }

            var columns = GetColumns(table);
            var rowsPerBatch = ResolveRowsPerBatch(totalRows, batchSize, columns.Length);

            SqliteCommand? command = null;
            var preparedRowsPerBatch = 0;
            try
            {
                for (var offset = 0; offset < totalRows; offset += rowsPerBatch)
                {
                    var currentRows = Math.Min(rowsPerBatch, totalRows - offset);
                    if (command == null || preparedRowsPerBatch != currentRows)
                    {
                        command?.Dispose();
                        command = CreatePreparedBulkInsertCommand(connection, activeTransaction ?? transaction, destinationTable, columns, currentRows);
                        preparedRowsPerBatch = currentRows;
                    }

                    ApplyBatchValues(command, columns, table, offset, currentRows);
                    ExecuteWithRetry(() => command.ExecuteNonQuery());
                }
            }
            finally
            {
                command?.Dispose();
            }

            if (!useTransaction)
            {
                transaction?.Commit();
            }
        }
        catch (DbaTransactionException)
        {
            throw;
        }
        catch (Exception ex)
        {
            if (!useTransaction)
            {
                transaction?.Rollback();
            }

            throw new DbaQueryExecutionException("Failed to execute bulk insert.", destinationTable, ex);
        }
        finally
        {
            if (!useTransaction)
            {
                transaction?.Dispose();
            }

            if (dispose)
            {
                connection?.Dispose();
            }
        }
    }

    /// <summary>
    /// Inserts all rows from the supplied <see cref="DataTable"/> into the specified destination table.
    /// </summary>
    public virtual void BulkInsert(
        string database,
        DataTable table,
        string destinationTable,
        bool useTransaction = false,
        int? batchSize = null)
        => BulkInsertCore(database, table, destinationTable, useTransaction, batchSize, useConnectionString: false);

    /// <summary>
    /// Inserts all rows from the supplied <see cref="DataTable"/> into the specified destination table using a SQLite connection string.
    /// </summary>
    public virtual void BulkInsertWithConnectionString(
        string connectionString,
        DataTable table,
        string destinationTable,
        bool useTransaction = false,
        int? batchSize = null)
        => BulkInsertCore(connectionString, table, destinationTable, useTransaction, batchSize, useConnectionString: true);

    private async Task BulkInsertCoreAsync(
        string database,
        DataTable table,
        string destinationTable,
        bool useTransaction,
        int? batchSize,
        CancellationToken cancellationToken,
        bool useConnectionString)
    {
        ValidateBulkInsertInputs(table, destinationTable, batchSize);

        var connectionString = ResolveBulkInsertConnectionString(database, useConnectionString);

        SqliteConnection? connection = null;
        var dispose = false;
        SqliteTransaction? activeTransaction = null;
        SqliteTransaction? transaction = null;
        try
        {
            (connection, activeTransaction, dispose) = await ResolveConnectionAsync(connectionString, useTransaction, cancellationToken).ConfigureAwait(false);

            if (!useTransaction)
            {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
                transaction = (SqliteTransaction)await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false);
#else
                transaction = connection.BeginTransaction();
#endif
            }

            var totalRows = table.Rows.Count;
            if (totalRows == 0)
            {
                return;
            }

            var columns = GetColumns(table);
            var rowsPerBatch = ResolveRowsPerBatch(totalRows, batchSize, columns.Length);

            SqliteCommand? command = null;
            var preparedRowsPerBatch = 0;
            try
            {
                for (var offset = 0; offset < totalRows; offset += rowsPerBatch)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    var currentRows = Math.Min(rowsPerBatch, totalRows - offset);
                    if (command == null || preparedRowsPerBatch != currentRows)
                    {
                        if (command != null)
                        {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
                            await command.DisposeAsync().ConfigureAwait(false);
#else
                            command.Dispose();
#endif
                        }

                        command = CreatePreparedBulkInsertCommand(connection, activeTransaction ?? transaction, destinationTable, columns, currentRows);
                        preparedRowsPerBatch = currentRows;
                    }

                    ApplyBatchValues(command, columns, table, offset, currentRows);
                    await ExecuteBulkInsertCommandAsync(command, cancellationToken).ConfigureAwait(false);
                }
            }
            finally
            {
                if (command != null)
                {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
                    await command.DisposeAsync().ConfigureAwait(false);
#else
                    command.Dispose();
#endif
                }
            }

            if (!useTransaction)
            {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
                if (transaction != null)
                {
                    await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
                }
#else
                transaction?.Commit();
#endif
            }
        }
        catch (DbaTransactionException)
        {
            throw;
        }
        catch (Exception ex)
        {
            if (!useTransaction)
            {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
                if (transaction != null)
                {
                    await RollbackOwnedBulkInsertTransactionAsync(transaction, CancellationToken.None).ConfigureAwait(false);
                }
#else
                transaction?.Rollback();
#endif
            }

            if (IsCallerCancellation(ex, cancellationToken))
            {
                throw;
            }

            throw CreateQueryExecutionOrCancellationException("Failed to execute bulk insert.", destinationTable, ex, cancellationToken);
        }
        finally
        {
            if (!useTransaction)
            {
                if (transaction != null)
                {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
                    await transaction.DisposeAsync().ConfigureAwait(false);
#else
                    transaction.Dispose();
#endif
                }
            }

            if (dispose)
            {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
                if (connection != null)
                {
                    await connection.DisposeAsync().ConfigureAwait(false);
                }
#else
                connection?.Dispose();
#endif
            }
        }
    }

    /// <summary>
    /// Asynchronously inserts all rows from the supplied <see cref="DataTable"/> into the specified destination table.
    /// </summary>
    public virtual Task BulkInsertAsync(
        string database,
        DataTable table,
        string destinationTable,
        bool useTransaction = false,
        int? batchSize = null,
        CancellationToken cancellationToken = default)
        => BulkInsertCoreAsync(database, table, destinationTable, useTransaction, batchSize, cancellationToken, useConnectionString: false);

    /// <summary>
    /// Asynchronously inserts all rows from the supplied <see cref="DataTable"/> into the specified destination table using a SQLite connection string.
    /// </summary>
    public virtual Task BulkInsertWithConnectionStringAsync(
        string connectionString,
        DataTable table,
        string destinationTable,
        bool useTransaction = false,
        int? batchSize = null,
        CancellationToken cancellationToken = default)
        => BulkInsertCoreAsync(connectionString, table, destinationTable, useTransaction, batchSize, cancellationToken, useConnectionString: true);

    private static DataColumn[] GetColumns(DataTable table)
    {
        var columns = new DataColumn[table.Columns.Count];
        for (var i = 0; i < table.Columns.Count; i++)
        {
            columns[i] = table.Columns[i];
        }

        return columns;
    }

    private static string ResolveBulkInsertConnectionString(string database, bool useConnectionString)
    {
        if (!useConnectionString)
        {
            return BuildOperationalConnectionString(database);
        }

        if (string.IsNullOrWhiteSpace(database))
        {
            throw new ArgumentException("Connection string cannot be null or whitespace.", nameof(database));
        }

        return NormalizeConnectionString(database);
    }

    internal static int ResolveRowsPerBatch(int totalRows, int? batchSize, int columnCount)
    {
        if (totalRows <= 0)
        {
            return 1;
        }

        if (columnCount <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(columnCount), "Column count must be greater than zero.");
        }

        if (columnCount > BulkInsertParameterLimit)
        {
            throw new ArgumentOutOfRangeException(nameof(columnCount), $"SQLite bulk insert supports at most {BulkInsertParameterLimit} columns per row.");
        }

        var parameterLimitedRows = Math.Max(1, BulkInsertParameterLimit / columnCount);
        if (batchSize.HasValue && batchSize.Value > 0)
        {
            return Math.Min(batchSize.Value, parameterLimitedRows);
        }

        return Math.Min(totalRows, Math.Min(DefaultBulkInsertBatchSize, parameterLimitedRows));
    }

    private static string BuildBulkInsertStatement(string destinationTable, DataColumn[] columns, int rowsPerBatch)
    {
        var columnNames = new string[columns.Length];
        for (var i = 0; i < columns.Length; i++)
        {
            columnNames[i] = QuoteIdentifier(columns[i].ColumnName);
        }

        var values = new string[rowsPerBatch];
        for (var rowIndex = 0; rowIndex < rowsPerBatch; rowIndex++)
        {
            var parameterNames = new string[columns.Length];
            for (var colIndex = 0; colIndex < columns.Length; colIndex++)
            {
                parameterNames[colIndex] = GetParameterName(rowIndex, colIndex);
            }

            values[rowIndex] = $"({string.Join(", ", parameterNames)})";
        }

        return $"INSERT INTO {QuoteIdentifierPath(destinationTable)} ({string.Join(", ", columnNames)}) VALUES {string.Join(", ", values)};";
    }

    private SqliteCommand CreatePreparedBulkInsertCommand(SqliteConnection connection, SqliteTransaction? transaction, string destinationTable, DataColumn[] columns, int rowsPerBatch)
    {
        var command = connection.CreateCommand();
        command.Transaction = transaction;
        command.CommandText = BuildBulkInsertStatement(destinationTable, columns, rowsPerBatch);
        ApplyCommandTimeout(command);

        for (var rowIndex = 0; rowIndex < rowsPerBatch; rowIndex++)
        {
            for (var colIndex = 0; colIndex < columns.Length; colIndex++)
            {
                command.Parameters.AddWithValue(GetParameterName(rowIndex, colIndex), DBNull.Value);
            }
        }

        command.Prepare();
        return command;
    }

    /// <summary>
    /// Executes the prepared SQLite bulk-insert command for the current batch.
    /// </summary>
    /// <param name="command">Prepared command populated with batch values.</param>
    /// <param name="cancellationToken">Token used to cancel execution and retry delays.</param>
    protected virtual Task ExecuteBulkInsertCommandAsync(SqliteCommand command, CancellationToken cancellationToken)
        => ExecuteWithRetryAsync(
            async () => {
                await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
                return 0;
            },
            cancellationToken);

    /// <summary>
    /// Rolls back an owned bulk-insert transaction during asynchronous failure cleanup.
    /// </summary>
    /// <param name="transaction">The owned transaction to roll back.</param>
    /// <param name="cancellationToken">Cleanup token used for rollback.</param>
    protected virtual Task RollbackOwnedBulkInsertTransactionAsync(SqliteTransaction transaction, CancellationToken cancellationToken)
    {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
        return transaction.RollbackAsync(cancellationToken);
#else
        transaction.Rollback();
        return Task.CompletedTask;
#endif
    }

    private static string GetParameterName(int rowIndex, int colIndex) => $"@p{rowIndex}_{colIndex}";

    private static string QuoteIdentifier(string identifier)
    {
        if (string.IsNullOrWhiteSpace(identifier))
        {
            throw new ArgumentException("Identifier cannot be null or whitespace.", nameof(identifier));
        }

        return "\"" + identifier.Replace("\"", "\"\"") + "\"";
    }

    private static string QuoteIdentifierPath(string identifierPath)
    {
        if (string.IsNullOrWhiteSpace(identifierPath))
        {
            throw new ArgumentException("Identifier cannot be null or whitespace.", nameof(identifierPath));
        }

        var parts = DbaIdentifierPath.SplitSegments(identifierPath);
        var quotedParts = new string[parts.Count];
        for (var i = 0; i < parts.Count; i++)
        {
            quotedParts[i] = QuoteIdentifier(DbaIdentifierPath.UnquoteSegment(parts[i]));
        }

        return string.Join(".", quotedParts);
    }

    private static void ApplyBatchValues(SqliteCommand command, DataColumn[] columns, DataTable table, int offset, int rowsPerBatch)
    {
        var parameterIndex = 0;
        for (var rowIndex = 0; rowIndex < rowsPerBatch; rowIndex++)
        {
            var row = table.Rows[offset + rowIndex];
            for (var colIndex = 0; colIndex < columns.Length; colIndex++)
            {
                command.Parameters[parameterIndex].Value = row[columns[colIndex]];
                parameterIndex++;
            }
        }
    }

    private static void ValidateBulkInsertInputs(DataTable table, string destinationTable, int? batchSize)
    {
        if (table == null)
        {
            throw new ArgumentNullException(nameof(table));
        }

        if (string.IsNullOrWhiteSpace(destinationTable))
        {
            throw new ArgumentException("Destination table cannot be null or whitespace.", nameof(destinationTable));
        }

        if (table.Columns.Count == 0)
        {
            throw new ArgumentException("Bulk insert requires at least one column.", nameof(table));
        }

        if (batchSize.HasValue && batchSize.Value <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(batchSize), "Batch size must be greater than zero.");
        }
    }
}