SQLite.AsyncCommandExecution.cs

using System;
using System.Collections.Generic;
using System.Data;
using System.Data.Common;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Data.Sqlite;

namespace DBAClientX;

public partial class SQLite
{
    /// <summary>
    /// Asynchronously executes a SQL query and materializes the result using the shared pipeline from <see cref="DatabaseClientBase"/>.
    /// </summary>
    public virtual async Task<object?> QueryAsync(
        string database,
        string query,
        IDictionary<string, object?>? parameters = null,
        bool useTransaction = false,
        CancellationToken cancellationToken = default,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null)
    {
        ValidateCommandText(query);
        var connectionString = BuildOperationalConnectionString(database);

        SqliteConnection? connection = null;
        SqliteTransaction? transaction = null;
        var dispose = false;
        try
        {
            (connection, transaction, dispose) = await ResolveConnectionAsync(connectionString, useTransaction, cancellationToken).ConfigureAwait(false);
            using var interrupt = RegisterOwnedStatementInterrupt(connection, dispose, cancellationToken, query);
            var dbTypes = ConvertParameterTypes(parameterTypes);
            return await ExecuteResolvedQueryAsync(connection, transaction, query, parameters, cancellationToken, dbTypes, parameterDirections).ConfigureAwait(false);
        }
        catch (Exception ex) when (!IsCallerCancellation(ex, cancellationToken))
        {
            throw CreateQueryExecutionOrCancellationException("Failed to execute query.", query, ex, cancellationToken);
        }
        finally
        {
            await DisposeOwnedResourceAsync(connection, dispose, DisposeSQLiteConnectionAsync).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Asynchronously executes a SQL query using a SQLite connection string and materializes the result using the shared pipeline.
    /// </summary>
    public virtual async Task<object?> QueryWithConnectionStringAsync(
        string connectionString,
        string query,
        IDictionary<string, object?>? parameters = null,
        bool useTransaction = false,
        CancellationToken cancellationToken = default,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null)
    {
        ValidateCommandText(query);
        var normalizedConnectionString = NormalizeConnectionString(connectionString);

        SqliteConnection? connection = null;
        SqliteTransaction? transaction = null;
        var dispose = false;
        try
        {
            (connection, transaction, dispose) = await ResolveConnectionAsync(normalizedConnectionString, useTransaction, cancellationToken).ConfigureAwait(false);
            using var interrupt = RegisterOwnedStatementInterrupt(connection, dispose, cancellationToken, query);
            var dbTypes = ConvertParameterTypes(parameterTypes);
            return await ExecuteResolvedQueryAsync(connection, transaction, query, parameters, cancellationToken, dbTypes, parameterDirections).ConfigureAwait(false);
        }
        catch (DbaTransactionException)
        {
            throw;
        }
        catch (Exception ex) when (!IsCallerCancellation(ex, cancellationToken) && (ex is DbException or InvalidOperationException or ArgumentException))
        {
            throw CreateQueryExecutionOrCancellationException("Failed to execute query.", query, ex, cancellationToken);
        }
        finally
        {
            await DisposeOwnedResourceAsync(connection, dispose, DisposeSQLiteConnectionAsync).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Asynchronously executes a SQL query and maps each returned row using the provided callback.
    /// </summary>
    public virtual Task<IReadOnlyList<T>> QueryAsync<T>(
        string database,
        string query,
        Func<IDataRecord, T> map,
        IDictionary<string, object?>? parameters = null,
        bool useTransaction = false,
        CancellationToken cancellationToken = default,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null)
        => QueryAsListAsync(database, query, map, parameters, useTransaction, cancellationToken, parameterTypes, parameterDirections);

    /// <summary>
    /// Asynchronously executes a SQL query and maps each returned row using the provided callback.
    /// </summary>
    public virtual async Task<IReadOnlyList<T>> QueryAsListAsync<T>(
        string database,
        string query,
        Func<IDataRecord, T> map,
        IDictionary<string, object?>? parameters = null,
        bool useTransaction = false,
        CancellationToken cancellationToken = default,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null,
        Action<IDataRecord>? initialize = null)
    {
        ValidateCommandText(query);
        if (map == null) throw new ArgumentNullException(nameof(map));

        var connectionString = BuildOperationalConnectionString(database);

        SqliteConnection? connection = null;
        SqliteTransaction? transaction = null;
        var dispose = false;
        try
        {
            (connection, transaction, dispose) = await ResolveConnectionAsync(connectionString, useTransaction, cancellationToken).ConfigureAwait(false);
            using var interrupt = RegisterOwnedStatementInterrupt(connection, dispose, cancellationToken, query);
            var dbTypes = ConvertParameterTypes(parameterTypes);
            return await ExecuteMappedQueryAsync(connection, transaction, query, map, initialize, parameters, cancellationToken, dbTypes, parameterDirections).ConfigureAwait(false);
        }
        catch (Exception ex) when (!IsCallerCancellation(ex, cancellationToken))
        {
            throw CreateQueryExecutionOrCancellationException("Failed to execute mapped query.", query, ex, cancellationToken);
        }
        finally
        {
            await DisposeOwnedResourceAsync(connection, dispose, DisposeSQLiteConnectionAsync).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Asynchronously executes a SQL query against a read-only connection and materializes the result using the shared pipeline from <see cref="DatabaseClientBase"/>.
    /// </summary>
    public virtual async Task<object?> QueryReadOnlyAsync(
        string database,
        string query,
        IDictionary<string, object?>? parameters = null,
        CancellationToken cancellationToken = default,
        int? busyTimeoutMs = null,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null)
    {
        ValidateCommandText(query);
        var connectionString = BuildOperationalConnectionString(database, readOnly: true);

        SqliteConnection? connection = null;
        var dispose = false;
        try
        {
            (connection, _, dispose) = await ResolveConnectionAsync(connectionString, useTransaction: false, cancellationToken, busyTimeoutMs).ConfigureAwait(false);
            using var interrupt = RegisterOwnedStatementInterrupt(connection, dispose, cancellationToken, query);
            var dbTypes = ConvertParameterTypes(parameterTypes);
            return await ExecuteResolvedQueryAsync(connection, null, query, parameters, cancellationToken, dbTypes, parameterDirections).ConfigureAwait(false);
        }
        catch (Exception ex) when (!IsCallerCancellation(ex, cancellationToken))
        {
            throw CreateQueryExecutionOrCancellationException("Failed to execute query.", query, ex, cancellationToken);
        }
        finally
        {
            await DisposeOwnedResourceAsync(connection, dispose, DisposeSQLiteConnectionAsync).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Asynchronously executes a SQL query against a read-only connection and maps each returned row using the provided callback.
    /// </summary>
    /// <typeparam name="T">The type produced by the row mapping callback.</typeparam>
    /// <param name="database">Path to the SQLite database file.</param>
    /// <param name="query">SQL query to execute.</param>
    /// <param name="map">Callback used to map the current <see cref="DbDataReader"/> row to <typeparamref name="T"/>.</param>
    /// <param name="parameters">Optional parameter values.</param>
    /// <param name="cancellationToken">Token used to cancel command execution.</param>
    /// <param name="busyTimeoutMs">Optional busy timeout in milliseconds.</param>
    /// <param name="parameterTypes">Optional provider-specific parameter types.</param>
    /// <param name="parameterDirections">Optional parameter directions.</param>
    /// <param name="initialize">
    /// Optional callback invoked once after the reader is opened and before the first row is read. Use it to cache ordinals.
    /// </param>
    /// <returns>A read-only list of mapped rows.</returns>
    public virtual async Task<IReadOnlyList<T>> QueryReadOnlyAsListAsync<T>(
        string database,
        string query,
        Func<DbDataReader, T> map,
        IDictionary<string, object?>? parameters = null,
        CancellationToken cancellationToken = default,
        int? busyTimeoutMs = null,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null,
        Action<DbDataReader>? initialize = null)
    {
        ValidateCommandText(query);
        if (map == null) throw new ArgumentNullException(nameof(map));

        var connectionString = BuildOperationalConnectionString(database, readOnly: true);

        SqliteConnection? connection = null;
        var dispose = false;
        try
        {
            (connection, _, dispose) = await ResolveConnectionAsync(connectionString, useTransaction: false, cancellationToken, busyTimeoutMs).ConfigureAwait(false);
            using var interrupt = RegisterOwnedStatementInterrupt(connection, dispose, cancellationToken, query);
            var dbTypes = ConvertParameterTypes(parameterTypes);

            var list = await ExecuteCommandWithDiagnosticsAsync(async () =>
            {
                using var command = connection.CreateCommand();
                command.CommandText = query;
                AddParameters(command, parameters, dbTypes, parameterDirections);
                ApplyCommandTimeout(command);

                using var reader = await AwaitWithCallerCancellationAsync(
                    () => command.ExecuteReaderAsync(CommandBehavior.SingleResult, cancellationToken),
                    cancellationToken).ConfigureAwait(false);
                initialize?.Invoke(reader);

                var results = new List<T>();
                while (await AwaitWithCallerCancellationAsync(
                    () => reader.ReadAsync(cancellationToken),
                    cancellationToken).ConfigureAwait(false))
                {
                    results.Add(map(reader));
                }

                UpdateOutputParameters(command, parameters);
                return (IReadOnlyList<T>)results;
            }, connection, transaction: null, query, "mapped", cancellationToken, static rows => rows.Count).ConfigureAwait(false);

            return list;
        }
        catch (OperationCanceledException ex) when (IsCallerCancellation(ex, cancellationToken))
        {
            throw;
        }
        catch (Exception ex) when (!IsCallerCancellation(ex, cancellationToken))
        {
            throw CreateQueryExecutionOrCancellationException("Failed to execute mapped query.", query, ex, cancellationToken);
        }
        finally
        {
            await DisposeOwnedResourceAsync(connection, dispose, DisposeSQLiteConnectionAsync).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Asynchronously executes a SQL statement that does not return rows.
    /// </summary>
    public virtual async Task<int> ExecuteNonQueryAsync(
        string database,
        string query,
        IDictionary<string, object?>? parameters = null,
        bool useTransaction = false,
        CancellationToken cancellationToken = default,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null)
    {
        ValidateCommandText(query);
        var connectionString = BuildOperationalConnectionString(database);
        return await ExecuteNonQueryCoreAsync(
            connectionString,
            query,
            parameters,
            useTransaction,
            cancellationToken,
            parameterTypes,
            parameterDirections).ConfigureAwait(false);
    }

    /// <summary>
    /// Asynchronously executes a SQL statement using a full SQLite connection string.
    /// </summary>
    public virtual async Task<int> ExecuteNonQueryWithConnectionStringAsync(
        string connectionString,
        string query,
        IDictionary<string, object?>? parameters = null,
        bool useTransaction = false,
        CancellationToken cancellationToken = default,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null)
    {
        ValidateCommandText(query);
        var normalizedConnectionString = NormalizeConnectionString(connectionString);
        return await ExecuteNonQueryCoreAsync(
            normalizedConnectionString,
            query,
            parameters,
            useTransaction,
            cancellationToken,
            parameterTypes,
            parameterDirections).ConfigureAwait(false);
    }

    private async Task<int> ExecuteNonQueryCoreAsync(
        string connectionString,
        string query,
        IDictionary<string, object?>? parameters,
        bool useTransaction,
        CancellationToken cancellationToken,
        IDictionary<string, SqliteType>? parameterTypes,
        IDictionary<string, ParameterDirection>? parameterDirections)
    {
        SqliteConnection? connection = null;
        SqliteTransaction? transaction = null;
        var dispose = false;
        try
        {
            (connection, transaction, dispose) = await ResolveConnectionAsync(connectionString, useTransaction, cancellationToken).ConfigureAwait(false);
            using var interrupt = RegisterOwnedStatementInterrupt(connection, dispose, cancellationToken, query);
            var dbTypes = ConvertParameterTypes(parameterTypes);
            return await base.ExecuteNonQueryAsync(connection, transaction, query, parameters, cancellationToken, dbTypes, parameterDirections).ConfigureAwait(false);
        }
        catch (Exception ex) when (!IsCallerCancellation(ex, cancellationToken))
        {
            throw CreateQueryExecutionOrCancellationException("Failed to execute non-query.", query, ex, cancellationToken);
        }
        finally
        {
            await DisposeOwnedResourceAsync(connection, dispose, DisposeSQLiteConnectionAsync).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Asynchronously executes a SQL query that returns a single scalar value.
    /// </summary>
    public virtual async Task<object?> ExecuteScalarAsync(
        string database,
        string query,
        IDictionary<string, object?>? parameters = null,
        bool useTransaction = false,
        CancellationToken cancellationToken = default,
        IDictionary<string, SqliteType>? parameterTypes = null,
        IDictionary<string, ParameterDirection>? parameterDirections = null)
    {
        ValidateCommandText(query);
        var connectionString = BuildOperationalConnectionString(database);

        SqliteConnection? connection = null;
        SqliteTransaction? transaction = null;
        var dispose = false;
        try
        {
            (connection, transaction, dispose) = await ResolveConnectionAsync(connectionString, useTransaction, cancellationToken).ConfigureAwait(false);
            using var interrupt = RegisterOwnedStatementInterrupt(connection, dispose, cancellationToken, query);
            var dbTypes = ConvertParameterTypes(parameterTypes);
            return await base.ExecuteScalarAsync(connection, transaction, query, parameters, cancellationToken, dbTypes, parameterDirections).ConfigureAwait(false);
        }
        catch (Exception ex) when (!IsCallerCancellation(ex, cancellationToken))
        {
            throw CreateQueryExecutionOrCancellationException("Failed to execute scalar query.", query, ex, cancellationToken);
        }
        finally
        {
            await DisposeOwnedResourceAsync(connection, dispose, DisposeSQLiteConnectionAsync).ConfigureAwait(false);
        }
    }

    private async Task<(SqliteConnection Connection, SqliteTransaction? Transaction, bool Dispose)> ResolveConnectionAsync(string connectionString, bool useTransaction, CancellationToken cancellationToken, int? busyTimeoutMs = null, bool retainTimeoutUpdates = false)
    {
        if (useTransaction)
        {
            lock (_syncRoot)
            {
                if (_transaction == null || _transactionConnection == null)
                {
                    throw new DbaTransactionException("Transaction has not been started.");
                }

                var normalizedConnectionString = NormalizeConnectionString(connectionString);
                if (_transactionConnectionString != null && !string.Equals(_transactionConnectionString, normalizedConnectionString, StringComparison.OrdinalIgnoreCase))
                {
                    throw new DbaTransactionException("The requested connection details do not match the active transaction.");
                }

                return (_transactionConnection, _transaction, false);
            }
        }

        return await SqliteTransientRetry.RunAsync(async retryToken =>
        {
            var connection = CreateConfiguredConnection(connectionString);
            try
            {
                await AwaitWithCallerCancellationAsync(
                    () => OpenConnectionWithDiagnosticsAsync(connection, retryToken),
                    retryToken).ConfigureAwait(false);
                await ApplyBusyTimeoutAsync(
                    connection,
                    ResolveConnectionBusyTimeout(connectionString, busyTimeoutMs),
                    retryToken).ConfigureAwait(false);
                ApplyConnectionConfiguration(connection, retainTimeoutUpdates: retainTimeoutUpdates);
                return (connection, (SqliteTransaction?)null, true);
            }
            catch
            {
                await DisposeSQLiteConnectionAsync(connection).ConfigureAwait(false);
                throw;
            }
        }, CreateTransientRetryOptions(), cancellationToken: cancellationToken).ConfigureAwait(false);
    }

    /// <summary>Executes a query after the SQLite connection and optional transaction have been resolved.</summary>
    protected virtual Task<object?> ExecuteResolvedQueryAsync(
        SqliteConnection connection,
        SqliteTransaction? transaction,
        string query,
        IDictionary<string, object?>? parameters,
        CancellationToken cancellationToken,
        IDictionary<string, DbType>? parameterTypes,
        IDictionary<string, ParameterDirection>? parameterDirections)
        => base.ExecuteQueryAsync(
            connection,
            transaction,
            query,
            parameters,
            cancellationToken,
            parameterTypes,
            parameterDirections);

    private static ValueTask DisposeSQLiteConnectionAsync(SqliteConnection connection)
    {
#if NETSTANDARD2_1_OR_GREATER || NETCOREAPP3_0_OR_GREATER || NET5_0_OR_GREATER
        return connection.DisposeAsync();
#else
        connection.Dispose();
        return default;
#endif
    }
}