sqlinsight-pro/src/SQLInsight.Collectors/Performance/QueryStoreCollector.cs

using System.Data.Common;
using SQLInsight.Core.Data;
using SQLInsight.Core.Instances;

namespace SQLInsight.Collectors.Performance;

/// <summary>
/// Reads Query Store from every user database that has it on. Only the 100 heaviest queries (by total duration over
/// eight days) are returned per database, so the cost stays bounded even with a large Query Store.
/// </summary>
public sealed class QueryStoreCollector : Collector<QueryStoreSnapshot>
{
    public const int TopQueries = 100;

    internal const string DatabasesSql = """
        SELECT name FROM sys.databases
        WHERE is_query_store_on = 1 AND state = 0 AND database_id > 4
        """;

    public override string Name => "query-store";
    public override TimeSpan DefaultInterval => TimeSpan.FromMinutes(30);
    public override IReadOnlyCollection<Capability> RequiredCapabilities => [Capability.QueryStore];

    protected override async Task<QueryStoreSnapshot> CollectAsync(CollectorContext context, QueryStoreSnapshot? previous)
    {
        var ct = context.CancellationToken;
        var names = await Db.QueryAsync(context.Connection, DatabasesSql, r => r.String(0), ct).ConfigureAwait(false);
        var databases = new List<QueryStoreDatabase>(names.Count);
        foreach (var name in names)
        {
            databases.Add(await ReadAsync(context.Connection, name, ct).ConfigureAwait(false));
        }

        return new QueryStoreSnapshot { CapturedAt = context.Now, Databases = databases };
    }

    private static async Task<QueryStoreDatabase> ReadAsync(DbConnection connection, string database, CancellationToken ct)
    {
        try
        {
            var plans = await Db.QueryAsync(connection, BuildStatsSql(database), r => new QueryPlanStats
            {
                QueryId = r.Long(0),
                PlanId = r.Long(1),
                IsRecent = r.Bool(2),
                Executions = r.Long(3),
                TotalDurationUs = r.Double(4),
                TotalCpuUs = r.Double(5),
                TotalLogicalReads = r.Double(6),
                MaxStdevDurationUs = r.Double(7),
                MaxDurationUs = r.Double(8),
                MinDurationUs = r.Double(9),
                LastExecution = r.IsNull(10) ? null : new DateTimeOffset(DateTime.SpecifyKind(r.DateTime(10), DateTimeKind.Utc)),
                IsForced = r.Bool(11),
                QueryText = r.NullableString(12),
                ObjectName = r.NullableString(13),
            }, ct).ConfigureAwait(false);

            var failed = await Db.QueryAsync(connection, BuildFailedForcedSql(database),
                r => new FailedForcedPlan(r.Long(0), r.Long(1), r.Long(2), r.NullableString(3), r.NullableString(4)), ct).ConfigureAwait(false);

            return new QueryStoreDatabase { Name = database, Plans = plans, FailedForcedPlans = failed };
        }
        catch (DbException ex)
        {
            return new QueryStoreDatabase { Name = database, Plans = [], Error = ex.Message };
        }
    }

    internal static string BuildStatsSql(string database)
    {
        var db = Db.Identifier(database);
        return $"""
            DECLARE @now datetimeoffset = SYSDATETIMEOFFSET();
            DECLARE @recent datetimeoffset = DATEADD(hour, -24, @now);
            DECLARE @history datetimeoffset = DATEADD(day, -8, @now);

            WITH stats AS (
                SELECT p.query_id, p.plan_id,
                       CASE WHEN i.start_time >= @recent THEN 1 ELSE 0 END AS is_recent,
                       SUM(rs.count_executions) AS execs,
                       SUM(rs.avg_duration * rs.count_executions) AS total_duration,
                       SUM(rs.avg_cpu_time * rs.count_executions) AS total_cpu,
                       SUM(rs.avg_logical_io_reads * rs.count_executions) AS total_reads,
                       MAX(rs.stdev_duration) AS max_stdev_duration,
                       MAX(rs.max_duration) AS max_duration,
                       MIN(rs.min_duration) AS min_duration,
                       MAX(rs.last_execution_time) AS last_execution
                FROM {db}.sys.query_store_runtime_stats rs
                JOIN {db}.sys.query_store_runtime_stats_interval i ON i.runtime_stats_interval_id = rs.runtime_stats_interval_id
                JOIN {db}.sys.query_store_plan p ON p.plan_id = rs.plan_id
                WHERE i.start_time >= @history
                GROUP BY p.query_id, p.plan_id, CASE WHEN i.start_time >= @recent THEN 1 ELSE 0 END
            ), top_queries AS (
                SELECT TOP ({TopQueries}) query_id FROM stats GROUP BY query_id ORDER BY SUM(total_duration) DESC
            )
            SELECT s.query_id, s.plan_id, s.is_recent, s.execs, s.total_duration, s.total_cpu, s.total_reads,
                   s.max_stdev_duration, s.max_duration, s.min_duration,
                   CAST(SWITCHOFFSET(s.last_execution, '+00:00') AS datetime2(0)),
                   p.is_forced_plan, LEFT(qt.query_sql_text, 2000),
                   OBJECT_SCHEMA_NAME(q.object_id, DB_ID({Db.Literal(database)})) + N'.' + OBJECT_NAME(q.object_id, DB_ID({Db.Literal(database)}))
            FROM stats s
            JOIN top_queries t ON t.query_id = s.query_id
            JOIN {db}.sys.query_store_plan p ON p.plan_id = s.plan_id
            JOIN {db}.sys.query_store_query q ON q.query_id = s.query_id
            JOIN {db}.sys.query_store_query_text qt ON qt.query_text_id = q.query_text_id;
            """;
    }

    internal static string BuildFailedForcedSql(string database)
    {
        var db = Db.Identifier(database);
        return $"""
            SELECT p.query_id, p.plan_id, p.force_failure_count, p.last_force_failure_reason_desc, LEFT(qt.query_sql_text, 2000)
            FROM {db}.sys.query_store_plan p
            JOIN {db}.sys.query_store_query q ON q.query_id = p.query_id
            JOIN {db}.sys.query_store_query_text qt ON qt.query_text_id = q.query_text_id
            WHERE p.is_forced_plan = 1 AND p.force_failure_count > 0;
            """;
    }
}