question
  → build prompt (schema as DDL + rules + examples)
  → model (IChatClient → Ollama)
  → parse reply (SQL | CANNOT_ANSWER)
  → validate (ScriptDom parse + allow-list)      ─┐
  → estimate cost (SET SHOWPLAN_XML ON)           ├─ on error: send it back to the model, max 2 rounds
  → execute as agent_reader (row cap, timeout)   ─┘
  → rows

dotnet add package Microsoft.Data.SqlClient
dotnet add package Microsoft.Extensions.AI
dotnet add package OllamaSharp
dotnet add package Microsoft.SqlServer.TransactSql.ScriptDom

CREATE ROLE agent_readers AUTHORIZATION dbo;
 
GRANT SELECT ON SCHEMA::sales     TO agent_readers;
GRANT SELECT ON SCHEMA::catalog   TO agent_readers;
GRANT SELECT ON SCHEMA::inventory TO agent_readers;
GRANT SELECT ON hr.Departments       TO agent_readers;
GRANT SELECT ON hr.EmployeeDirectory TO agent_readers;  -- a view of hr.Employees without Salary
DENY  SELECT ON hr.Employees         TO agent_readers;
GRANT SHOWPLAN TO agent_readers;                         -- for the cost check in step 6
 
CREATE USER agent_reader WITHOUT LOGIN WITH DEFAULT_SCHEMA = sales;
ALTER ROLE agent_readers ADD MEMBER agent_reader;

using System.Data;
using Microsoft.Data.SqlClient;
 
internal sealed class ImpersonatedSession : IAsyncDisposable
{
    private readonly byte[] _cookie;
 
    private ImpersonatedSession(SqlConnection connection, byte[] cookie)
    {
        Connection = connection;
        _cookie = cookie;
    }
 
    public SqlConnection Connection { get; }
 
    public static async Task<ImpersonatedSession> OpenAsync(string connectionString, string agentUser, CancellationToken ct)
    {
        var connection = new SqlConnection(connectionString);
        try
        {
            await connection.OpenAsync(ct);
            await using var cmd = connection.CreateCommand();
            // Ad hoc batch on purpose: WITH COOKIE fails inside sp_executesql (error 15590).
            // agentUser comes from configuration and is validated as a simple identifier.
            cmd.CommandText = $"""
                DECLARE @cookie varbinary(8000);
                EXECUTE AS USER = N'{agentUser}' WITH COOKIE INTO @cookie;
                SELECT @cookie;
                """;
            var cookie = await cmd.ExecuteScalarAsync(ct) as byte[]
                ?? throw new InvalidOperationException("EXECUTE AS did not return a cookie.");
            return new ImpersonatedSession(connection, cookie);
        }
        catch
        {
            await connection.DisposeAsync();
            throw;
        }
    }
 
    public async ValueTask DisposeAsync()
    {
        try
        {
            if (Connection.State == ConnectionState.Open)
            {
                await using var cmd = Connection.CreateCommand();
                cmd.CommandText = $"REVERT WITH COOKIE = 0x{Convert.ToHexString(_cookie)};";
                await cmd.ExecuteNonQueryAsync();
            }
        }
        catch (Exception ex) when (ex is SqlException or InvalidOperationException)
        {
            // Never hand a still-impersonated connection back to the pool.
            SqlConnection.ClearPool(Connection);
        }
        finally
        {
            await Connection.DisposeAsync();
        }
    }
}

public async Task<QueryResult> ExecuteAsync(string sql, int maxRows, int timeoutSeconds, CancellationToken ct)
{
    await using var session = await ImpersonatedSession.OpenAsync(_connectionString, "agent_reader", ct);
    await using var cmd = session.Connection.CreateCommand();
    cmd.CommandText = sql;
    cmd.CommandTimeout = timeoutSeconds;
    await using var reader = await cmd.ExecuteReaderAsync(ct);
 
    var columns = Enumerable.Range(0, reader.FieldCount).Select(reader.GetName).ToList();
    var rows = new List<object?[]>();
    var truncated = false;
    while (await reader.ReadAsync(ct))
    {
        if (rows.Count == maxRows)
        {
            truncated = true;
            cmd.Cancel(); // stop the server producing rows nobody will read
            break;
        }
        var row = new object?[reader.FieldCount];
        for (var i = 0; i < row.Length; i++)
            row[i] = reader.IsDBNull(i) ? null : reader.GetValue(i);
        rows.Add(row);
    }
    return new QueryResult(columns, rows, truncated);
}

SELECT s.name AS SchemaName, o.name AS ObjectName, c.name AS ColumnName,
       ty.name AS TypeName, c.is_nullable,
       CAST(ep.value AS nvarchar(1000)) AS Description
FROM sys.columns c
JOIN sys.objects o ON o.object_id = c.object_id
JOIN sys.schemas s ON s.schema_id = o.schema_id
JOIN sys.types ty ON ty.user_type_id = c.user_type_id
LEFT JOIN sys.extended_properties ep
       ON ep.class = 1 AND ep.major_id = c.object_id
      AND ep.minor_id = c.column_id AND ep.name = N'MS_Description'
WHERE o.type IN ('U', 'V') AND o.is_ms_shipped = 0
ORDER BY s.name, o.name, c.column_id;

-- Sales orders (header).
TABLE sales.Orders (
  OrderId int PRIMARY KEY,
  CustomerId int NOT NULL -> sales.Customers.CustomerId,
  OrderDate datetime2 NOT NULL,
  Status varchar(10) NOT NULL -- values: 'Pending', 'Shipped', 'Delivered', 'Cancelled', 'Returned'
)

public const string SystemPrompt = """
    You are an expert Microsoft SQL Server developer. You write one read-only T-SQL query that answers
    the user's question using ONLY the tables and columns in the schema provided.
 
    Rules:
    - T-SQL only: use TOP (n), never LIMIT; use GETDATE(), DATEADD, DATEDIFF, DATEFROMPARTS, EOMONTH, YEAR(), MONTH().
    - Always write schema-qualified table names (e.g. sales.Orders) and give tables short aliases.
    - Join only along the foreign keys shown (the "->" arrows). Never guess a column that is not listed.
    - For enum-like columns, compare only with the listed values, spelled exactly.
    - Read the descriptions: they define business meaning (for example how revenue is calculated).
    - One SELECT statement (common table expressions allowed). Never modify data.
    - If the question needs data that is not in the schema, or asks to change data, do not write a query:
      reply with exactly: CANNOT_ANSWER: <one short reason>
 
    Reply format: one short sentence saying what the query returns, then the query in a ```sql block.
    """;

var messages = new List<ChatMessage> { new(ChatRole.System, SystemPrompt + "\n\n" + context) };
 
foreach (var ex in examples)
{
    messages.Add(new ChatMessage(ChatRole.User, ex.Question));
    messages.Add(new ChatMessage(ChatRole.Assistant, $"{ex.Explanation}\n```sql\n{ex.Sql}\n```"));
}
 
messages.Add(new ChatMessage(ChatRole.User, question));

builder.Services.AddChatClient(sp =>
    {
        // 127.0.0.1, not localhost: on Windows localhost resolves to ::1 first and Ollama listens
        // on IPv4 only, so every new connection waited ~2 s for the IPv6 attempt to fail.
        var http = new HttpClient
        {
            BaseAddress = new Uri("http://127.0.0.1:11434"),
            Timeout = TimeSpan.FromSeconds(180), // the first call also loads the model
        };
        return new OllamaApiClient(http, "qwen2.5-coder:7b");
    })
    .Use(inner => new ConfigureOptionsChatClient(inner,
        // Ollama's default context is 2048 tokens: too small for a schema-aware prompt.
        o => o.AddOllamaOption(OllamaOption.NumCtx, 8192)))
    .UseLogging();

public sealed record ParsedResponse(string? Sql, string? Explanation, string? CannotAnswerReason);
 
public static partial class ModelResponseParser
{
    [GeneratedRegex(@"```[ \t]*(?:sql|tsql|t-sql|mssql)?[ \t]*\r?\n(?<code>.*?)```",
        RegexOptions.Singleline | RegexOptions.IgnoreCase)]
    private static partial Regex FencedBlock();
 
    [GeneratedRegex(@"CANNOT_ANSWER\s*[:\-]?\s*(?<reason>[^\r\n]*)")]
    private static partial Regex CannotAnswer();
 
    [GeneratedRegex(@"^\s*(WITH|SELECT)\b", RegexOptions.IgnoreCase | RegexOptions.Multiline)]
    private static partial Regex QueryStart();
 
    public static ParsedResponse Parse(string? response)
    {
        var text = (response ?? "").Trim();
 
        var refusal = CannotAnswer().Match(text);
        if (refusal.Success && !FencedBlock().IsMatch(text))
            return new ParsedResponse(null, null, refusal.Groups["reason"].Value.Trim());
 
        var fence = FencedBlock().Match(text);
        if (fence.Success)
            return new ParsedResponse(fence.Groups["code"].Value.Trim().TrimEnd(';'),
                text.Remove(fence.Index, fence.Length).Trim(), null);
 
        // No fence: take everything from the first line that starts a query.
        var start = QueryStart().Match(text);
        return start.Success
            ? new ParsedResponse(text[start.Index..].Trim().TrimEnd(';'), text[..start.Index].Trim(), null)
            : new ParsedResponse(null, null, null);
    }
}

using Microsoft.SqlServer.TransactSql.ScriptDom;
 
public static class SqlValidator
{
    public static IReadOnlyList<string> Validate(string sql, IReadOnlySet<string> allowedTables)
    {
        var parser = new TSql170Parser(initialQuotedIdentifiers: true);
        var fragment = parser.Parse(new StringReader(sql), out var errors);
        if (errors.Count > 0)
            return errors.Take(3).Select(e => $"Syntax error at line {e.Line}, column {e.Column}: {e.Message}").ToList();
 
        var statements = ((TSqlScript)fragment).Batches.SelectMany(b => b.Statements).ToList();
        if (statements is not [SelectStatement select])
            return ["Exactly one SELECT statement is allowed."];
        if (select.Into is not null)
            return ["SELECT ... INTO creates a table and is not allowed."];
 
        var visitor = new AllowListVisitor(allowedTables);
        select.Accept(visitor);
        return visitor.Problems;
    }
 
    private sealed class AllowListVisitor(IReadOnlySet<string> allowedTables) : TSqlFragmentVisitor
    {
        private readonly HashSet<string> _ctes = new(StringComparer.OrdinalIgnoreCase);
 
        public List<string> Problems { get; } = [];
 
        public override void ExplicitVisit(SelectStatement node)
        {
            // Register CTE names first so references to them aren't reported as unknown tables.
            foreach (var cte in node.WithCtesAndXmlNamespaces?.CommonTableExpressions ?? [])
                _ctes.Add(cte.ExpressionName.Value);
            base.ExplicitVisit(node);
        }
 
        public override void Visit(TableReference node)
        {
            if (node is not (NamedTableReference or QueryDerivedTable or JoinTableReference or JoinParenthesisTableReference))
                Problems.Add($"Not allowed: table source {node.GetType().Name}."); // OPENROWSET, OPENQUERY, TVFs, ...
        }
 
        public override void Visit(NamedTableReference node)
        {
            var name = node.SchemaObject;
            if (name.ServerIdentifier is not null || name.DatabaseIdentifier is not null)
            {
                Problems.Add("Not allowed: references to other databases or linked servers.");
                return;
            }
            if (name.SchemaIdentifier is null && _ctes.Contains(name.BaseIdentifier.Value))
                return;
            var qualified = $"{name.SchemaIdentifier?.Value}.{name.BaseIdentifier.Value}";
            if (!allowedTables.Contains(qualified))
                Problems.Add($"Table '{qualified}' does not exist or is not available. Available tables: {string.Join(", ", allowedTables)}.");
        }
 
        public override void Visit(VariableReference node) =>
            Problems.Add($"Not allowed: variable {node.Name}.");
 
        public override void Visit(FunctionCall node)
        {
            if (node.CallTarget is not null)
                Problems.Add($"Not allowed: user-defined function '{node.FunctionName.Value}'.");
        }
    }
}

public async Task<double> EstimateCostAsync(string sql, CancellationToken ct)
{
    await using var session = await ImpersonatedSession.OpenAsync(_connectionString, "agent_reader", ct);
    await using (var on = session.Connection.CreateCommand())
    {
        on.CommandText = "SET SHOWPLAN_XML ON;";
        await on.ExecuteNonQueryAsync(ct);
    }
    try
    {
        await using var cmd = session.Connection.CreateCommand();
        cmd.CommandText = sql;
        var plan = XDocument.Parse((string)(await cmd.ExecuteScalarAsync(ct))!);
        return plan.Descendants()
            .Where(e => e.Name.LocalName == "StmtSimple")
            .Sum(e => double.Parse(e.Attribute("StatementSubTreeCost")?.Value ?? "0", CultureInfo.InvariantCulture));
    }
    finally
    {
        // Must be switched off before the session reverts: under SHOWPLAN nothing executes, REVERT included.
        await using var off = session.Connection.CreateCommand();
        off.CommandText = "SET SHOWPLAN_XML OFF;";
        await off.ExecuteNonQueryAsync(CancellationToken.None);
    }
}

private static void AddRepairTurn(List<ChatMessage> messages, string reply, IReadOnlyList<string> problems)
{
    var sb = new StringBuilder("That query was not accepted:\n");
    foreach (var p in problems)
        sb.Append("- ").Append(p).Append('\n');
    sb.Append("\nWrite a corrected query using only tables and columns from the schema. ")
      .Append("If a column was taken from the wrong table, take it from the table that has it (joining it if needed). ")
      .Append("Reply with CANNOT_ANSWER: <reason> only if the data needed does not exist anywhere in the schema.");
 
    messages.Add(new ChatMessage(ChatRole.Assistant, reply));
    messages.Add(new ChatMessage(ChatRole.User, sb.ToString()));
}

builder.WebHost.ConfigureKestrel(k => k.Limits.MaxRequestBodySize = 64 * 1024);
 
builder.Services.AddRateLimiter(o =>
{
    o.RejectionStatusCode = StatusCodes.Status429TooManyRequests;
    o.AddPolicy("ask", http => RateLimitPartition.GetFixedWindowLimiter(
        http.Connection.RemoteIpAddress?.ToString() ?? "unknown",
        _ => new FixedWindowRateLimiterOptions { PermitLimit = 20, Window = TimeSpan.FromMinutes(1) }));
});
 
var app = builder.Build();
app.UseRateLimiter();
 
app.MapPost("/api/ask", async (AskRequest request, TextToSqlAgent agent, CancellationToken ct) =>
    {
        if (string.IsNullOrWhiteSpace(request.Question))
            return Results.Problem("Ask a question.", statusCode: StatusCodes.Status400BadRequest);
        try
        {
            return Results.Ok(await agent.AskAsync(request.Question, ct));
        }
        catch (HttpRequestException)
        {
            return Results.Problem("The language model is not reachable. Is Ollama running?",
                statusCode: StatusCodes.Status503ServiceUnavailable);
        }
    })
    .RequireRateLimiting("ask");
 
public sealed record AskRequest(string Question);