任意のSELECT文から SELECT COUNT(*) FROM (SELECT ～ ) を動的に生成する (ScriptDomを使ったSELECT文の書き換え)

こちらの記事が大変参考になりました(というよりもこちらの記事を見て思い付きました)。
T-SQL の パーサーとジェネレーター(ScriptDom)の紹介 - お だ のスペース
ScriptDom でクエリを改変する - お だ のスペース

動機

ページング処理用に検索結果の総レコード数を取得するため、検索 SQL を文字列置換して SELECT COUNT(*) FROM ([元の SQL]) に書き換えようと頑張っているコードを見かけました。
しかし、単純に SQL を COUNT(*) のクエリで包むだけでは不十分で、内側の SQL から ORDER BY 句を取り除く必要があります。（一部の場合を除きインラインビュー内には ORDER BY 句が書けない）
これらの SQL 文の操作において、単純な文字列置換では SQL コメントを含む場合などに対応しきれなくなる可能性があります。
可能であれば SQL パーサーを使用した正確な SQL 書き換えを行なうべきでしょう。

方法

Microsoft® SQL Server® 2012 Feature Pack で提供される Microsoft.SqlServer.TransactSql.ScriptDom を使用し、
入力となる SQL を文法的に正しく解析して ORDER BY 句の除去＋ SELECT COUNT(*) FROM ([元の SQL]) の形に書き換え、までを行ないます。

ちなみに生成された SQL には表別名指定の AS が入っちゃうので Oracle では使えません。
（AS 句の出力も SqlScriptGeneratorOptions で調整出来るようになるといいですね）

コード

using System;
using System.Collections.Generic;
using System.IO;
using Microsoft.SqlServer.TransactSql.ScriptDom;

public class RowCountRewriteVisitor : TSqlFragmentVisitor
{
    private SelectStatement CreateSelectStatement()
    {
        // 内側の select 1 (=>QueryDerivedTable) を件数取得したいクエリに置き換える
        var template = @"select count(*) as cnt from ( select 1 ) tbl";
        IList<ParseError> ignored;
        var parsed = (TSqlScript)new TSql110Parser(false).Parse(new StringReader(template), out ignored);
        return (SelectStatement)parsed.Batches[0].Statements[0];
    }

    public override void Visit(TSqlBatch node)
    {
        for (int i = 0; i < node.Statements.Count; i++)
        {
            if (node.Statements[i].GetType() == typeof(SelectStatement))
            {
                var originalStatement = (SelectStatement)node.Statements[i];
                // テンプレートになるクエリ (SelectStatement) を生成して元のクエリをテンプレートクエリの内側に入れる
                var newStatement = this.CreateSelectStatement();
                var counting = (QuerySpecification)newStatement.QueryExpression;
                var innerStatement = (QueryDerivedTable)counting.FromClause.TableReferences[0];
                innerStatement.QueryExpression = originalStatement.QueryExpression;
                // 内側のクエリに ORDER BY 句は要らない
                innerStatement.QueryExpression.OrderByClause = null;

                node.Statements[i] = newStatement;
            }
        }
        base.Visit(node);
    }
}

public class EnumerateVisitor : TSqlFragmentVisitor
{
    private string query;

    public EnumerateVisitor(string query)
    {
        this.query = query;
    }

    public override void Visit(TSqlFragment fragment)
    {
        Console.WriteLine(fragment.ToString() + ": " + this.query.Substring(fragment.StartOffset, fragment.FragmentLength));
        base.Visit(fragment);
    }
}

呼び出し側

var query = @"select null order by 1; select 1 as value1, 2 as value2 from tbl union select 3 as value1, 4 as value2 from tbl  order by 1; select 99";

var parser = new TSql110Parser(false);
IList<ParseError> errors;
var parsed = parser.Parse(new StringReader(query), out errors);
if (errors.Count != 0) { throw new Exception("Invalid SQL!"); }

var v = new RowCountRewriteVisitor();
parsed.Accept(v);

string formatted;
new Sql110ScriptGenerator().GenerateScript(parsed, out formatted);
formatted;

こういうコード片を VisualStudio で簡単に実行する方法が良く分からなかったので CShell を利用。
1. Microsoft.SqlServer.TransactSql.ScriptDom を参照に追加
2. RowCountRewriteVisitor クラスをロード
3. 呼び出し側のコードを実行
で OK。

CShell 便利です。

実行結果

SELECT count(*) AS cnt
FROM   (SELECT NULL) AS tbl;

SELECT count(*) AS cnt
FROM   (SELECT 1 AS value1,
               2 AS value2
        FROM   tbl
        UNION
        SELECT 3 AS value1,
               4 AS value2
        FROM   tbl) AS tbl;

SELECT count(*) AS cnt
FROM   (SELECT 99) AS tbl;

複数のSQLステートメントがある場合でもちゃんと書き換えできました。

課題

アプリケーション側でもSQLのハードパースのような処理を検索実行時に毎回させることになるので、パフォーマンスがかなりアレな予感です。

あらかた書いてから気付きましたが、元 SQL が SELECT TOP n ～ だった場合は恐らくその SQL から ORDER BY 句を削除すべきではないでしょう。上記の実装にはこのあたりの考慮が入っていません。
（QuerySpecification#TopRowFilter とか怪しそうですね）

〆

実戦投入どころか大した検証もしていません。
ご利用の際は自己責任でお願いします。

補足

Q. QuerySpecification やら QueryDerivedTable やらはどこから出てきたのか？
A. EnumerateVisitor で中身見て MSDN 漁って目星を付けました。(とてもだるい)

Microsoft.SqlServer.TransactSql.ScriptDom.TSqlScript: select count(*) as cnt from (select 3 as value3, 4 as value4 from tbl2) tbl
Microsoft.SqlServer.TransactSql.ScriptDom.TSqlBatch: select count(*) as cnt from (select 3 as value3, 4 as value4 from tbl2) tbl
Microsoft.SqlServer.TransactSql.ScriptDom.SelectStatement: select count(*) as cnt from (select 3 as value3, 4 as value4 from tbl2) tbl
Microsoft.SqlServer.TransactSql.ScriptDom.QuerySpecification: select count(*) as cnt from (select 3 as value3, 4 as value4 from tbl2) tbl
Microsoft.SqlServer.TransactSql.ScriptDom.SelectScalarExpression: count(*) as cnt
Microsoft.SqlServer.TransactSql.ScriptDom.FunctionCall: count(*)
Microsoft.SqlServer.TransactSql.ScriptDom.Identifier: count
Microsoft.SqlServer.TransactSql.ScriptDom.ColumnReferenceExpression: *
Microsoft.SqlServer.TransactSql.ScriptDom.IdentifierOrValueExpression: cnt
Microsoft.SqlServer.TransactSql.ScriptDom.Identifier: cnt
Microsoft.SqlServer.TransactSql.ScriptDom.FromClause: from (select 3 as value3, 4 as value4 from tbl2) tbl
Microsoft.SqlServer.TransactSql.ScriptDom.QueryDerivedTable: (select 3 as value3, 4 as value4 from tbl2) tbl
Microsoft.SqlServer.TransactSql.ScriptDom.Identifier: tbl
Microsoft.SqlServer.TransactSql.ScriptDom.QuerySpecification: select 3 as value3, 4 as value4 from tbl2
Microsoft.SqlServer.TransactSql.ScriptDom.SelectScalarExpression: 3 as value3
Microsoft.SqlServer.TransactSql.ScriptDom.IntegerLiteral: 3
Microsoft.SqlServer.TransactSql.ScriptDom.IdentifierOrValueExpression: value3
Microsoft.SqlServer.TransactSql.ScriptDom.Identifier: value3
Microsoft.SqlServer.TransactSql.ScriptDom.SelectScalarExpression: 4 as value4
Microsoft.SqlServer.TransactSql.ScriptDom.IntegerLiteral: 4
Microsoft.SqlServer.TransactSql.ScriptDom.IdentifierOrValueExpression: value4
Microsoft.SqlServer.TransactSql.ScriptDom.Identifier: value4
Microsoft.SqlServer.TransactSql.ScriptDom.FromClause: from tbl2
Microsoft.SqlServer.TransactSql.ScriptDom.NamedTableReference: tbl2
Microsoft.SqlServer.TransactSql.ScriptDom.SchemaObjectName: tbl2
Microsoft.SqlServer.TransactSql.ScriptDom.Identifier: tbl2