Add profiling and throughput reporting to the benchmarks

benchmarks/README.md

@@ -1,8 +1,10 @@
 # Benchmarks
 
 Scaling benchmarks for the code paths that have carried quadratic-CPU
-(DoS) regressions.  They are developer tools, not part of the test suite:
-they are not run by `pytest` and not wired into CI.
+(DoS) regressions, plus `bench_parse_throughput.py`, which measures
+ordinary large SQL rather than an advisory shape.  They are developer
+tools, not part of the test suite: they are not run by `pytest` and not
+wired into CI.
 
 Run one, or all of them at once:
 
@@ -17,18 +19,32 @@ Every script accepts the same options:
 --sizes 100,200,400     comma-separated sizes, overriding the defaults
 --timeout SECONDS       abort a single measurement (default: 30)
 --vector SUBSTRING      only run matching vectors, repeatable
+--profile               profile instead of timing, see below
+--profile-top N         frames to print per vector (default: 15)
 ```
 
+`--profile` runs each vector once at its largest size under `cProfile` and
+prints the hottest frames by `tottime`, the time spent in the frame itself.
+That is the column that names the function to change; `cumtime` mostly
+re-reports the callers above it.  Combine with `--vector` to profile a
+single path:
+
+```
+python benchmarks/bench_parse_throughput.py --profile --vector UNION
+```
+
+Profiling makes no scaling claim, so `--profile` always exits 0.
+
 ## Output
 
 Each vector prints one row per size and closes with a verdict:
 
 ```
 wide column lists:
-         n      input         time    ratio  status
-       500     4402 B      12.3 ms        -  ok
-      1000     8902 B      24.6 ms    2.00x  ok
-      2000    18.9 kB      49.1 ms    1.99x  ok
+         n      input         time    ratio     kB/s  status
+       500     4402 B      12.3 ms        -      358  ok
+      1000     8902 B      24.6 ms    2.00x      362  ok
+      2000    18.9 kB      49.1 ms    1.99x      385  ok
   scaling exponent 1.00 (1.0 linear, 2.0 quadratic)  =>  linear
 ```
 
@@ -39,6 +55,24 @@ behaviour.  A script exits non-zero as soon as one vector grows
 super-linearly, which makes it usable as a regression check for the
 advisories it covers.
 
+`kB/s` is throughput at that size and does not feed the verdict.  It
+covers the blind spot growth alone leaves: a change that makes everything
+uniformly twice as slow keeps the exponent at 1.0 and is invisible in the
+`ratio` column, but halves this one.  Compare it against a run of the same
+script on the previous commit, not against the numbers above.
+
+Two caveats when reading a verdict:
+
+- The exponent is diluted by whatever linear work shares the path.  A
+  quadratic step behind a large linear one (lexing, typically) can fit
+  under the 1.5 threshold at small sizes and still be quadratic; the fit
+  uses only the largest sizes for that reason.  If the `ratio` column
+  climbs steadily above 2.00x while the verdict says linear, trust the
+  ratios and extend `--sizes`.
+- A vector sized to stay under `MAX_GROUPING_TOKENS` measures what the
+  default limits permit, not whether the path is linear.  The cap bounds
+  how far a quadratic path can be pushed; it does not straighten it.
+
 The `status` column is `ok`, `cap` when a grouping limit rejected the
 input before the measured path was reached, or `timeout`.  Only `ok` rows
 above 1 ms enter the fit; a vector without at least two of them is

benchmarks/_harness.py

@@ -25,16 +25,26 @@ def wide_select(n):
 """
 
 import argparse
+import cProfile
+import io
 import math
+import pstats
 import signal
 import time
 from collections.abc import Callable
 from dataclasses import dataclass
+from pathlib import Path
 
 import sqlparse
 
 DEFAULT_TIMEOUT = 30
 
+#: Frames printed per vector by --profile.
+DEFAULT_PROFILE_TOP = 15
+
+#: Stripped from profile output so frames read as "sqlparse/engine/...".
+_REPO_ROOT = f'{Path(sqlparse.__file__).resolve().parent.parent}/'
+
 #: Midpoint between linear (1.0) and quadratic (2.0) growth.  At or above
 #: this the measured code path is reported as super-linear.
 SUPERLINEAR_EXPONENT = 1.5
@@ -134,16 +144,19 @@ def _exponent(points):
 def run_vector(vector, timeout=DEFAULT_TIMEOUT):
     """Measure one vector, print its table and return its verdict."""
     print(f'{vector.name}:')
-    print(f'  {"n":>8}  {"input":>9}  {"time":>11}  {"ratio":>7}  status')
+    print(f'  {"n":>8}  {"input":>9}  {"time":>11}  {"ratio":>7}'
+          f'  {"kB/s":>7}  status')
 
     points = []
     previous = None
     for n in vector.sizes:
         sql = vector.build(n)
         status, elapsed = measure(vector.run, sql, timeout)
         ratio = f'{elapsed / previous:.2f}x' if previous else '-'
+        rate = (f'{len(sql) / elapsed:.0f}'
+                if status == 'ok' and elapsed else '-')
         print(f'  {n:>8}  {_format_size(len(sql)):>9}  {elapsed:>8.1f} ms  '
-              f'{ratio:>7}  {status}')
+              f'{ratio:>7}  {rate:>7}  {status}')
         previous = elapsed
         if status == 'ok' and elapsed >= MIN_SIGNIFICANT_MS:
             points.append((n, elapsed))
@@ -159,6 +172,47 @@ def run_vector(vector, timeout=DEFAULT_TIMEOUT):
     return verdict
 
 
+def profile_vector(vector, top=DEFAULT_PROFILE_TOP, timeout=DEFAULT_TIMEOUT):
+    """Profile ``vector`` at its largest size, printing the hottest frames.
+
+    Sorted by ``tottime``, excluding callees: that points at the
+    function to change rather than at whatever calls it.
+    """
+    n = vector.sizes[-1]
+    sql = vector.build(n)
+    print(f'{vector.name}:  n={n}, {_format_size(len(sql))}')
+
+    profiler = cProfile.Profile()
+    status = 'ok'
+    if _CAN_ALARM:
+        signal.alarm(timeout)
+    profiler.enable()
+    try:
+        vector.run(sql)
+    except sqlparse.exceptions.SQLParseError:
+        status = 'cap'
+    except Timeout:
+        status = 'timeout'
+    finally:
+        profiler.disable()
+        if _CAN_ALARM:
+            signal.alarm(0)
+
+    if status != 'ok':
+        print(f'  {status}: the measured path was not reached, '
+              f'nothing to profile\n')
+        return
+
+    buffer = io.StringIO()
+    stats = pstats.Stats(profiler, stream=buffer)
+    stats.sort_stats('tottime').print_stats(top)
+    text = buffer.getvalue().replace(_REPO_ROOT, '')
+    header = text.find('ncalls')
+    for line in text[header if header > 0 else 0:].splitlines():
+        print(f'  {line}' if line.strip() else '')
+    print()
+
+
 def _parse_args(argv, title, vectors):
     parser = argparse.ArgumentParser(
         description=f'sqlparse benchmark: {title}',
@@ -173,6 +227,13 @@ def _parse_args(argv, title, vectors):
         '--vector', action='append', metavar='SUBSTRING',
         help='only run vectors whose name contains SUBSTRING '
              f'(available: {", ".join(v.name for v in vectors)})')
+    parser.add_argument(
+        '--profile', action='store_true',
+        help='instead of timing, profile each vector at its largest size '
+             'and print the hottest frames')
+    parser.add_argument(
+        '--profile-top', type=int, default=DEFAULT_PROFILE_TOP,
+        metavar='N', help='frames to print per vector with --profile')
     return parser.parse_args(argv)
 
 
@@ -196,6 +257,16 @@ def main(title, vectors, note=None, argv=None):
         print(note)
     print()
 
+    if args.profile:
+        print('profiling at the largest size of each vector, '
+              'sorted by tottime\n')
+        for vector in vectors:
+            profile_vector(
+                vector if args.sizes is None
+                else Vector(vector.name, vector.build, args.sizes, vector.run),
+                args.profile_top, args.timeout)
+        return 0  # hotspots only; no scaling claim, so nothing to fail on
+
     verdicts = [
         run_vector(
             vector if args.sizes is None

benchmarks/bench_parse_throughput.py

@@ -0,0 +1,100 @@
+"""Parse throughput benchmark for realistic large SQL.
+
+The other benchmarks here measure pathological shapes tied to specific
+advisories; this one measures the ordinary case, the large varied SQL that
+BI tools and ETL jobs generate.  A plain slowdown shows up in the ``kB/s``
+column rather than in the exponent -- a uniform 2x regression leaves the
+exponent at 1.0.
+
+Run with:  python benchmarks/bench_parse_throughput.py
+           python benchmarks/bench_parse_throughput.py --profile
+"""
+
+import sys
+
+from _harness import Vector, main
+
+import sqlparse
+from sqlparse.engine import grouping
+
+# Disable the grouping-stage DoS guards.  They are a policy limit on
+# untrusted input, not a property of the engine: realistic SQL of the size
+# measured here exceeds MAX_GROUPING_TOKENS and would be rejected outright.
+grouping.MAX_GROUPING_DEPTH = None
+grouping.MAX_GROUPING_TOKENS = None
+
+
+def analytics_query(n):
+    """One CTE selecting n expressions, each with a CASE and a window
+    function, over a 7-way join, with a long IN list and a GROUP BY."""
+    projections = ',\n'.join(
+        f"""        t{i % 7}.col_{i} AS alias_{i},
+        CASE WHEN t{i % 7}.flag_{i} = 'Y' AND t{i % 7}.amt_{i} > {i}.5
+             THEN COALESCE(t{i % 7}.amt_{i} * 1.075, 0)
+             ELSE NULL END AS calc_{i},
+        SUM(t{i % 7}.amt_{i}) OVER (
+            PARTITION BY t{i % 7}.grp_{i} ORDER BY t{i % 7}.dt_{i}
+            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS win_{i}"""
+        for i in range(n))
+    joins = ''.join(
+        f"    LEFT JOIN schema_a.table_{i} t{i} "
+        f"ON t{i}.id = t0.id AND t{i}.dt >= DATE '2024-01-01'\n"
+        for i in range(1, 7))
+    in_list = ', '.join(f"'st_{i}'" for i in range(n))
+    group_by = ', '.join(f't0.grp_{i}' for i in range(max(n // 4, 1)))
+    return (
+        f'WITH base AS (\n    SELECT\n{projections}\n'
+        f'    FROM schema_a.table_0 t0\n{joins}'
+        f'    WHERE t0.status IN ({in_list})\n'
+        f"      AND t0.region = 'EMEA'  -- restrict region\n"
+        f'    GROUP BY {group_by}\n'
+        f'    HAVING COUNT(*) > 10\n)\n'
+        f'SELECT * FROM base ORDER BY alias_0 DESC')
+
+
+def statement_script(n):
+    """A migration-style script of n separate statements."""
+    return '\n'.join(
+        f"UPDATE tbl_{i % 20} SET col_a = 'v{i}', col_b = {i} "
+        f"WHERE id = {i} AND flag = 'Y';"
+        for i in range(n))
+
+
+def union_all_chain(n):
+    """n branches unioned together, the shape generated pivots produce."""
+    return '\nUNION ALL\n'.join(
+        f"SELECT {i} AS bucket, col_a, col_b FROM src_{i % 12} "
+        f"WHERE dt = DATE '2024-01-{i % 28 + 1:02d}' AND col_c > {i}"
+        for i in range(n))
+
+
+def string_and_comment_heavy(n):
+    """Comment lines, escaped string literals and inline comments."""
+    return '\n'.join(
+        f'-- comment line {i} explaining the next bit\n'
+        f"SELECT 'literal {i} with ''escaped'' quotes', /* inline {i} */ "
+        f"col_{i} FROM t WHERE x = '{i}';"
+        for i in range(n))
+
+
+def reindent(sql):
+    return sqlparse.format(sql, reindent=True)
+
+
+VECTORS = [
+    Vector('analytics query (CTE, joins, windows)',
+           analytics_query, (25, 50, 100, 200)),
+    Vector('multi-statement script',
+           statement_script, (100, 200, 400, 800)),
+    Vector('UNION ALL chain',
+           union_all_chain, (100, 200, 400, 800)),
+    Vector('string and comment heavy',
+           string_and_comment_heavy, (150, 300, 600, 1200)),
+    Vector('analytics query, format(reindent=True)',
+           analytics_query, (25, 50, 100, 200), reindent),
+]
+
+if __name__ == '__main__':
+    sys.exit(main('parse throughput', VECTORS,
+                  note='realistic large SQL -- watch the kB/s column',
+                  argv=sys.argv[1:]))

sqlparse/utils.py

@@ -5,6 +5,7 @@
 # This module is part of python-sqlparse and is released under
 # the BSD License: https://opensource.org/licenses/BSD-3-Clause
 
+import functools
 import itertools
 import re
 from collections import deque
@@ -67,6 +68,7 @@ def recurse(*cls):
     :return: function
     """
     def wrap(f):
+        @functools.wraps(f)
         def wrapped_f(tlist):
             for sgroup in tlist.get_sublists():
                 if not isinstance(sgroup, cls):