migcheck 0.1.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
migcheck/__init__.py ADDED
File without changes
migcheck/__main__.py ADDED
@@ -0,0 +1,240 @@
1
+ """Checks that a release's migrations don't break the release before it.
2
+
3
+ During a rolling deploy, or after a rollback, the previous release runs
4
+ against the new schema. migcheck records the queries the previous release
5
+ makes, then replays them against its schema and against the new one. A query
6
+ that works on the old schema and fails on the new one breaks production.
7
+ """
8
+
9
+ import argparse
10
+ import os
11
+ import shutil
12
+ import subprocess
13
+ import sys
14
+ import tempfile
15
+ import time
16
+ from collections import defaultdict
17
+
18
+ from .mysql import MySQL
19
+ from .postgres import Postgres
20
+
21
+ ADAPTERS = {"postgres": Postgres, "mysql": MySQL}
22
+ # Marks migcheck's pull request comment, so a later run can update it.
23
+ MARKER = "<!-- migcheck -->"
24
+ MARKDOWN = None
25
+
26
+
27
+ def main():
28
+ parser = argparse.ArgumentParser(prog="migcheck", description=(__doc__ or "").splitlines()[0])
29
+ parser.add_argument("--repo", default=".", help="git repository to check (default: .)")
30
+ parser.add_argument("--base", required=True, help="ref of the release running in production")
31
+ parser.add_argument("--head", required=True, help="ref of the release being deployed")
32
+ parser.add_argument("--migrate", required=True, help="shell command that applies migrations to $DATABASE_URL")
33
+ parser.add_argument("--workload", required=True, help="shell command that exercises the app against $DATABASE_URL, such as its test suite")
34
+ parser.add_argument("--db", choices=ADAPTERS, default="postgres", help="database the app uses (default: postgres)")
35
+ defaults = " or ".join(a.default_image for a in ADAPTERS.values())
36
+ parser.add_argument("--image", help=f"database Docker image (default: {defaults})")
37
+ parser.add_argument("--head-name", help="name to show for --head in the result, such as \"this pull request\"")
38
+ parser.add_argument("--markdown", metavar="PATH", help="also write the result as Markdown to PATH, for a pull request comment")
39
+ args = parser.parse_args()
40
+ global MARKDOWN
41
+ MARKDOWN = args.markdown
42
+
43
+ repo = os.path.abspath(args.repo)
44
+ workdir = tempfile.mkdtemp(prefix="migcheck-")
45
+ base_dir = os.path.join(workdir, "base")
46
+ head_dir = os.path.join(workdir, "head")
47
+ adapter = ADAPTERS[args.db]
48
+ db = adapter(args.image or adapter.default_image)
49
+ try:
50
+ _git(repo, "worktree", "add", "--detach", "-q", base_dir, args.base)
51
+ _git(repo, "worktree", "add", "--detach", "-q", head_dir, args.head)
52
+
53
+ step(f"Starting {db.image}")
54
+ db.start()
55
+
56
+ step(f"Applying {args.base} migrations")
57
+ db.create_db("mc_template")
58
+ _sh(args.migrate, base_dir, db.url("mc_template"))
59
+
60
+ step(f"Recording {args.base} workload queries")
61
+ db.create_db("mc_capture", template="mc_template")
62
+ db.record()
63
+ # Failing tests still issue queries, so a failing workload is a warning.
64
+ workload_exit = _sh(args.workload, base_dir, db.url("mc_capture"), check=False)
65
+ statements = db.statements(exclude={"mc_template"})
66
+ if not statements:
67
+ fail("the workload ran no queries, so there is nothing to check")
68
+
69
+ step(f"Applying {args.head} migrations")
70
+ db.create_db("mc_old", template="mc_template")
71
+ db.create_db("mc_new", template="mc_template")
72
+ _sh(args.migrate, head_dir, db.url("mc_new"))
73
+
74
+ step(f"Replaying {len(statements)} statements against both schemas")
75
+ old = db.replay("mc_old", statements)
76
+ new = db.replay("mc_new", statements)
77
+ except Exception as e: # Exit code 1 means a query breaks, so nothing else may use it.
78
+ fail(f"{type(e).__name__}: {e}")
79
+ finally:
80
+ db.stop()
81
+ for d in (base_dir, head_dir):
82
+ if os.path.isdir(d):
83
+ _git(repo, "worktree", "remove", "--force", d)
84
+ shutil.rmtree(workdir, ignore_errors=True)
85
+
86
+ sys.exit(report(db, args.base, args.head, statements, old, new, workload_exit, args.head_name))
87
+
88
+
89
+ def report(db, base, head, statements, old, new, workload_exit, head_name=None):
90
+ """Prints the result, writes it as Markdown if requested, and returns the exit code."""
91
+ head_md = head_name or f"`{head}`"
92
+ head = head_name or head
93
+ skipped = sum(1 for state, _ in old if state)
94
+ if skipped == len(statements):
95
+ fail(
96
+ f"all {skipped} statements fail on {base}'s own schema, so there is nothing to compare. "
97
+ f"First error: [{old[0][0]}] {old[0][1]}"
98
+ )
99
+ roots, knock_on = _classify(db, statements, old, new)
100
+ count = sum(len(v) for v in roots.values())
101
+ note = ""
102
+ if workload_exit:
103
+ note = f"The workload exited with code {workload_exit}. If it stopped early, this result covers only the queries it ran."
104
+
105
+ print()
106
+ if skipped:
107
+ print(f"Skipped {skipped} statements that also fail on {base}'s schema.")
108
+ if not roots:
109
+ print(f"PASS: all {len(statements) - skipped} statements from {base} still work after {head}'s migrations.")
110
+ else:
111
+ print(f"FAIL: {count} statements from {base} break after {head}'s migrations.\n")
112
+ for (state, msg), stmts in roots.items():
113
+ print(f" [{state}] {msg}")
114
+ print(f" {len(stmts)} statements in {_blocks(stmts)} transactions. Examples:")
115
+ for sql in _examples(stmts, 160):
116
+ print(f" {sql}")
117
+ print()
118
+ if knock_on:
119
+ print(" Later failures in the same transactions, likely caused by the errors above:")
120
+ for (state, msg), n in knock_on.items():
121
+ print(f" [{state}] {msg} ({n})")
122
+ print()
123
+ print(f"{base} fails during a rolling deploy of {head}, and after a rollback to {base}.")
124
+ if note:
125
+ print(f"\nNote: {note}")
126
+
127
+ if MARKDOWN:
128
+ md = [MARKER]
129
+ if not roots:
130
+ md.append("### migcheck: PASS")
131
+ md.append("")
132
+ md.append(f"All {len(statements) - skipped:,} statements from `{base}` still work after {head_md}'s migrations.")
133
+ else:
134
+ md.append("### migcheck: FAIL")
135
+ md.append("")
136
+ md.append(
137
+ f"{count:,} statements from `{base}` break after {head_md}'s migrations. "
138
+ f"`{base}` fails during a rolling deploy of {head_md}, and after a rollback to `{base}`."
139
+ )
140
+ for (state, msg), stmts in roots.items():
141
+ md.append("")
142
+ md.append(f"**`{state}` {_md_escape(msg)}**")
143
+ md.append(f"{len(stmts):,} statements in {_blocks(stmts):,} transactions. Examples:")
144
+ md.append("")
145
+ md.append("```sql")
146
+ md.extend(_examples(stmts, 500))
147
+ md.append("```")
148
+ if knock_on:
149
+ md.append("")
150
+ md.append("<details><summary>Later failures in the same transactions, likely caused by the errors above</summary>")
151
+ md.append("")
152
+ for (state, msg), n in knock_on.items():
153
+ md.append(f"- `{state}` {_md_escape(msg)} ({n:,})")
154
+ md.append("")
155
+ md.append("</details>")
156
+ if skipped:
157
+ md.append("")
158
+ md.append(f"Skipped {skipped:,} statements that also fail on `{base}`'s schema.")
159
+ if note:
160
+ md.append("")
161
+ md.append(f"> [!NOTE]\n> {note}")
162
+ _write_markdown(md)
163
+ return 1 if roots else 0
164
+
165
+
166
+ def _classify(db, statements, old, new):
167
+ """Splits the statements that break into root causes and knock-on errors.
168
+
169
+ The first break in a transaction is a root cause. A later integrity error
170
+ in the same transaction usually follows from it, such as a foreign key to
171
+ a row whose insert failed. A foreign key error after any earlier break is
172
+ also a knock-on, because apps often commit each request separately, so the
173
+ failed insert and the reference to it can sit in different transactions.
174
+ """
175
+ roots = defaultdict(list)
176
+ knock_on = defaultdict(int)
177
+ broken_blocks = set()
178
+ for s, (old_state, _), (new_state, new_msg) in zip(statements, old, new):
179
+ if old_state is not None or new_state is None:
180
+ continue
181
+ key = (new_state, new_msg)
182
+ follows_break = s.block in broken_blocks and db.is_integrity_error(new_state)
183
+ if key not in roots and (follows_break or (broken_blocks and db.is_foreign_key_error(new_state))):
184
+ knock_on[key] += 1
185
+ continue
186
+ broken_blocks.add(s.block)
187
+ roots[key].append(s)
188
+ return roots, knock_on
189
+
190
+
191
+ def _blocks(stmts):
192
+ return len({s.block for s in stmts})
193
+
194
+
195
+ def _examples(stmts, width):
196
+ examples = dict.fromkeys(" ".join(s.sql.split())[:width] for s in stmts)
197
+ return list(examples)[:3]
198
+
199
+
200
+ def _md_escape(text):
201
+ return text.replace("<", "&lt;").replace("*", "\\*").replace("_", "\\_").replace("`", "'")
202
+
203
+
204
+ def _write_markdown(lines):
205
+ with open(MARKDOWN, "w") as f:
206
+ f.write("\n".join(lines) + "\n")
207
+
208
+
209
+ def fail(msg):
210
+ """Exits with code 2, which means migcheck couldn't compare."""
211
+ print(f"migcheck: {msg}", file=sys.stderr)
212
+ if MARKDOWN:
213
+ _write_markdown([MARKER, "### migcheck: couldn't compare", "", "```", msg, "```"])
214
+ sys.exit(2)
215
+
216
+
217
+ def step(msg):
218
+ print(f"==> {time.strftime('%H:%M:%S')} {msg}", file=sys.stderr)
219
+
220
+
221
+ def _git(repo, *args):
222
+ subprocess.run(["git", "-C", repo, *args], check=True)
223
+
224
+
225
+ def _sh(cmd, cwd, url, check=True):
226
+ proc = subprocess.run(
227
+ cmd, shell=True, cwd=cwd, env={**os.environ, "DATABASE_URL": url},
228
+ capture_output=True, text=True,
229
+ )
230
+ if proc.returncode != 0 and not check:
231
+ print(f"migcheck: warning: command exited with {proc.returncode}: {cmd}", file=sys.stderr)
232
+ sys.stderr.write(proc.stderr[-2000:])
233
+ elif proc.returncode != 0:
234
+ sys.stderr.write(proc.stdout + proc.stderr)
235
+ fail(f"command failed in {cwd}: {cmd}")
236
+ return proc.returncode
237
+
238
+
239
+ if __name__ == "__main__":
240
+ main()
migcheck/common.py ADDED
@@ -0,0 +1,23 @@
1
+ """Types and helpers that the database adapters share."""
2
+
3
+ import subprocess
4
+ from dataclasses import dataclass, field
5
+
6
+
7
+ @dataclass
8
+ class Statement:
9
+ sql: str
10
+ # Statements from one workload transaction share a block. Replay rolls
11
+ # back a block's effects if the workload rolled back that transaction,
12
+ # the way a test runner discards each test's data.
13
+ block: int
14
+ keep: bool
15
+ # [table, id] for each row the statement inserted, in insert order.
16
+ ids: list = field(default_factory=list)
17
+
18
+
19
+ def run(cmd, stdin=None):
20
+ proc = subprocess.run(cmd, input=stdin, capture_output=True, text=True)
21
+ if proc.returncode != 0:
22
+ raise RuntimeError(f"{' '.join(cmd[:4])}... failed:\n{proc.stderr.strip()}")
23
+ return proc.stdout
migcheck/mysql.py ADDED
@@ -0,0 +1,338 @@
1
+ """MySQL adapter.
2
+
3
+ Runs a throwaway MySQL server in Docker, records the workload with the
4
+ general query log, and talks to the server with the mysql client inside the
5
+ container, so the host needs only Docker.
6
+ """
7
+
8
+ import bisect
9
+ import re
10
+ from collections import Counter
11
+ import subprocess
12
+ import time
13
+ import uuid
14
+
15
+ from .common import Statement
16
+ from .common import run as _run
17
+
18
+ DEFAULT_IMAGE = "mysql:8.4"
19
+ PASSWORD = "migcheck"
20
+ LOG_FILE = "/tmp/general.log"
21
+ SYSTEM_DBS = {"mysql", "information_schema", "performance_schema", "sys"}
22
+
23
+ # A general log entry starts with a timestamp, the connection ID, and the
24
+ # command. A statement's own newlines continue it on the following lines.
25
+ ENTRY = re.compile(rb"^\d{4}-\d\d-\d\dT[\d:.]+Z?\t\s*(\d+) ([A-Za-z ]+?)\t(.*)$", re.DOTALL)
26
+ CONNECT_DB = re.compile(r" on (\S+) using ")
27
+
28
+ TXN_BEGIN = re.compile(r"^\s*(BEGIN|START\s+TRANSACTION)\b(?!\s*;?\s*\S)", re.IGNORECASE)
29
+ TXN_COMMIT = re.compile(r"^\s*COMMIT\b", re.IGNORECASE)
30
+ TXN_ROLLBACK = re.compile(r"^\s*ROLLBACK\b(?!\s+(WORK\s+)?TO\b)", re.IGNORECASE)
31
+ AUTOCOMMIT = re.compile(
32
+ r"^\s*SET\s+(?:SESSION\s+|@@(?:SESSION\.)?)?autocommit\s*=\s*(\S+?)\s*;?\s*$", re.IGNORECASE
33
+ )
34
+ USE = re.compile(r"^\s*USE\s+`?([^`\s;]+)`?\s*;?\s*$", re.IGNORECASE)
35
+ # Statements that only make sense in the client's own session. Replay issues
36
+ # its own transaction control. SQL-level PREPARE and EXECUTE are dropped
37
+ # because the general log also records the statement each EXECUTE runs.
38
+ SESSION_ONLY = re.compile(
39
+ r"^\s*(XA|LOCK\s+TABLES?|UNLOCK\s+TABLES?|PREPARE|EXECUTE|DEALLOCATE|SET\s+TRANSACTION)\b", re.IGNORECASE
40
+ )
41
+ # Schema changes come from test runners building their database, not from the
42
+ # app's runtime queries. They commit any open transaction.
43
+ DDL = re.compile(r"^\s*(CREATE|ALTER|DROP|RENAME|GRANT|REVOKE)\b", re.IGNORECASE)
44
+ # Test runners empty tables between tests by name. On the new schema a listed
45
+ # table can be gone, or a new table can reference one and block the TRUNCATE,
46
+ # so replay empties the table only if it exists, ignoring foreign keys.
47
+ TRUNCATE = re.compile(r"^\s*TRUNCATE\s+(?:TABLE\s+)?((?:`[^`]+`|[\w$]+)(?:\.(?:`[^`]+`|[\w$]+))?)\s*;?\s*$", re.IGNORECASE)
48
+ ERROR = re.compile(r"^ERROR (\d+) \(\w+\) at line (\d+): (.*)$")
49
+
50
+ INTEGRITY_ERRORS = {"1022", "1048", "1062", "1169", "1216", "1217", "1364", "1451", "1452", "1557", "1586", "3819"}
51
+ FOREIGN_KEY_ERRORS = {"1216", "1217", "1451", "1452"}
52
+
53
+ REPLAY_SETUP = """\
54
+ DROP PROCEDURE IF EXISTS _migcheck_truncate;
55
+ DELIMITER //
56
+ CREATE PROCEDURE _migcheck_truncate(name varchar(64))
57
+ BEGIN
58
+ IF EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = name) THEN
59
+ SET @_migcheck_fk = @@foreign_key_checks, foreign_key_checks = 0;
60
+ SET @_migcheck_sql = CONCAT('TRUNCATE TABLE `', REPLACE(name, '`', '``'), '`');
61
+ PREPARE _migcheck_stmt FROM @_migcheck_sql;
62
+ EXECUTE _migcheck_stmt;
63
+ DEALLOCATE PREPARE _migcheck_stmt;
64
+ SET foreign_key_checks = @_migcheck_fk;
65
+ END IF;
66
+ END//
67
+ DELIMITER ;
68
+ """
69
+
70
+
71
+ class MySQL:
72
+ name = "MySQL"
73
+ default_image = DEFAULT_IMAGE
74
+
75
+ def __init__(self, image=DEFAULT_IMAGE):
76
+ self.image = image
77
+ self.container = None
78
+ self.port = None
79
+
80
+ def start(self):
81
+ self.container = f"migcheck-{uuid.uuid4().hex[:8]}"
82
+ _run([
83
+ "docker", "run", "-d", "--rm", "--name", self.container,
84
+ "-e", f"MYSQL_ROOT_PASSWORD={PASSWORD}",
85
+ "-p", "127.0.0.1::3306",
86
+ self.image,
87
+ f"--general-log-file={LOG_FILE}",
88
+ "--max-allowed-packet=1G",
89
+ "--skip-log-bin",
90
+ "--skip-name-resolve",
91
+ "--innodb-flush-log-at-trx-commit=0",
92
+ "--innodb-doublewrite=OFF",
93
+ ])
94
+ mapping = _run(["docker", "port", self.container, "3306/tcp"])
95
+ self.port = int(mapping.splitlines()[0].rsplit(":", 1)[1])
96
+ # The image's init step runs a socket-only server first, so a TCP
97
+ # check waits for the real one.
98
+ for _ in range(120):
99
+ check = subprocess.run(self._client("-h", "127.0.0.1", "-e", "SELECT 1"), capture_output=True)
100
+ if check.returncode == 0:
101
+ return
102
+ time.sleep(0.5)
103
+ raise RuntimeError("MySQL did not become ready in 60 seconds")
104
+
105
+ def stop(self):
106
+ if self.container:
107
+ subprocess.run(["docker", "rm", "-f", self.container], capture_output=True)
108
+ self.container = None
109
+
110
+ def url(self, db):
111
+ return f"mysql://root:{PASSWORD}@127.0.0.1:{self.port}/{db}"
112
+
113
+ def create_db(self, name, template=None):
114
+ """Creates a database. MySQL has no template databases, so a
115
+ template's tables and rows are copied with mysqldump."""
116
+ _run(self._client("-e", f"CREATE DATABASE `{name}`"))
117
+ if template:
118
+ _run(["docker", "exec", "-e", f"MYSQL_PWD={PASSWORD}", self.container, "sh", "-c",
119
+ "set -o pipefail; mysqldump -uroot --single-transaction --routines --triggers --events "
120
+ f"--hex-blob --no-tablespaces --set-gtid-purged=OFF {template} | mysql -uroot {name}"])
121
+
122
+ def record(self):
123
+ """Starts recording statements. The log stays off until the workload
124
+ runs, so setup statements don't appear in it."""
125
+ _run(self._client("-e", "SET GLOBAL general_log = ON"))
126
+
127
+ @staticmethod
128
+ def is_integrity_error(state):
129
+ return state in INTEGRITY_ERRORS
130
+
131
+ @staticmethod
132
+ def is_foreign_key_error(state):
133
+ return state in FOREIGN_KEY_ERRORS
134
+
135
+ def statements(self, exclude):
136
+ """Returns the statements the workload ran against any database not in
137
+ `exclude`.
138
+
139
+ Test runners often create their own database, such as Django's
140
+ test_<name>, so this doesn't filter on a single name.
141
+ """
142
+ self._stop_recording()
143
+ # The log can be far larger than the statements kept from it, so read
144
+ # it as a stream instead of loading it whole.
145
+ assert self.container
146
+ proc = subprocess.Popen(["docker", "exec", self.container, "cat", LOG_FILE], stdout=subprocess.PIPE)
147
+ assert proc.stdout is not None
148
+ statements = parse_log(proc.stdout, exclude)
149
+ if proc.wait() != 0:
150
+ raise RuntimeError(f"reading the MySQL log failed with exit code {proc.returncode}")
151
+ return statements
152
+
153
+ def replay(self, db, statements):
154
+ """Runs each statement in order and returns a list of (error code, message).
155
+
156
+ A statement that succeeds keeps its effects until the end of its
157
+ block, so later statements see the rows earlier ones wrote. MySQL
158
+ rolls back a failing statement on its own, and replay continues. Both
159
+ entries are None on success.
160
+ """
161
+ script, starts = build_script(statements)
162
+ assert self.container
163
+ proc = subprocess.run(
164
+ ["docker", "exec", "-i", "-e", f"MYSQL_PWD={PASSWORD}", self.container, "sh", "-c",
165
+ f"mysql -uroot --force --binary-mode --max-allowed-packet=1G --default-character-set=utf8mb4 {db} >/dev/null"],
166
+ input=script, capture_output=True,
167
+ )
168
+ results = parse_errors(proc.stderr.decode("utf-8", "replace"), starts, len(statements))
169
+ if proc.returncode not in (0, 1):
170
+ raise RuntimeError(f"mysql exited with code {proc.returncode}:\n{proc.stderr.decode(errors='replace')[-2000:]}")
171
+ return results
172
+
173
+ def _stop_recording(self):
174
+ # The log records this statement before it takes effect. Connecting
175
+ # to a system database keeps it out of the workload, because
176
+ # parse_log skips everything sent to those.
177
+ _run(self._client("-D", "mysql", "-e", "SET GLOBAL general_log = OFF"))
178
+
179
+ def _client(self, *args):
180
+ assert self.container
181
+ return ["docker", "exec", "-e", f"MYSQL_PWD={PASSWORD}", self.container, "mysql", "-uroot", *args]
182
+
183
+
184
+ class _Session:
185
+ def __init__(self):
186
+ self.db = None
187
+ self.autocommit = True
188
+ self.block = None # the open transaction's block, if any
189
+
190
+
191
+ def parse_log(log, exclude):
192
+ """Returns the statements in a general query log, read from the binary
193
+ line stream `log`, that ran against any database not in `exclude`."""
194
+ exclude = SYSTEM_DBS | set(exclude)
195
+ blocks = [] # [keep, [sql, ...]] in the order transactions started
196
+ sessions = {} # connection ID -> _Session
197
+
198
+ def close(session, keep):
199
+ if session.block is not None:
200
+ session.block[0] = keep
201
+ session.block = None
202
+
203
+ def add(session, sql):
204
+ if session.db in exclude:
205
+ return
206
+ if session.block is None and not session.autocommit:
207
+ session.block = [True, []]
208
+ blocks.append(session.block)
209
+ if session.block is not None:
210
+ session.block[1].append(sql)
211
+ else:
212
+ blocks.append([True, [sql]])
213
+
214
+ for conn, command, arg in _entries(log):
215
+ session = sessions.setdefault(conn, _Session())
216
+ if command == "Connect":
217
+ m = CONNECT_DB.search(arg)
218
+ session.db = m.group(1) if m else None
219
+ elif command == "Init DB":
220
+ session.db = arg.strip()
221
+ elif command == "Quit":
222
+ # A transaction still open at disconnect never committed.
223
+ close(sessions.pop(conn), False)
224
+ elif command in ("Query", "Execute"):
225
+ sql = arg.strip()
226
+ if m := AUTOCOMMIT.match(sql):
227
+ session.autocommit = m.group(1).strip("'\"").upper() in ("1", "ON", "TRUE")
228
+ if session.autocommit:
229
+ close(session, True)
230
+ elif TXN_BEGIN.match(sql):
231
+ close(session, True)
232
+ session.block = [True, []]
233
+ blocks.append(session.block)
234
+ elif TXN_COMMIT.match(sql):
235
+ close(session, True)
236
+ elif TXN_ROLLBACK.match(sql):
237
+ close(session, False)
238
+ elif m := USE.match(sql):
239
+ session.db = m.group(1)
240
+ elif DDL.match(sql):
241
+ close(session, True)
242
+ elif SESSION_ONLY.match(sql):
243
+ pass
244
+ elif m := TRUNCATE.match(sql):
245
+ # TRUNCATE commits any open transaction and can't be undone.
246
+ close(session, True)
247
+ autocommit, session.autocommit = session.autocommit, True
248
+ add(session, _lenient_truncate(m.group(1)))
249
+ session.autocommit = autocommit
250
+ else:
251
+ add(session, sql)
252
+ # A transaction still open when the workload exits never committed.
253
+ for session in sessions.values():
254
+ close(session, False)
255
+ return [
256
+ Statement(sql, i, keep)
257
+ for i, (keep, sqls) in enumerate(b for b in blocks if b[1])
258
+ for sql in sqls
259
+ ]
260
+
261
+
262
+ def build_script(statements):
263
+ """Returns the replay script for `statements` and the script line each
264
+ statement starts on."""
265
+ parts = [REPLAY_SETUP]
266
+ line = 1 + REPLAY_SETUP.count("\n")
267
+ starts = [] # first script line of each statement
268
+ sizes = Counter(s.block for s in statements)
269
+ current = None
270
+ in_txn = False
271
+ end_txn = ""
272
+ for s in statements:
273
+ if s.block != current:
274
+ if in_txn:
275
+ parts.append(end_txn)
276
+ line += 1
277
+ current = s.block
278
+ # A block with one statement runs on its own, as the
279
+ # workload's autocommit statements did.
280
+ in_txn = sizes[s.block] > 1 or not s.keep
281
+ end_txn = "COMMIT;\n" if s.keep else "ROLLBACK;\n"
282
+ if in_txn:
283
+ parts.append("START TRANSACTION;\n")
284
+ line += 1
285
+ starts.append(line)
286
+ # The terminator goes on its own line so a trailing comment can't
287
+ # swallow it.
288
+ text = s.sql.rstrip().rstrip(";") + "\n;\n"
289
+ parts.append(text)
290
+ line += text.count("\n")
291
+ if in_txn:
292
+ parts.append(end_txn)
293
+ return "".join(parts).encode("utf-8", "surrogateescape"), starts
294
+
295
+
296
+ def parse_errors(stderr, starts, count):
297
+ """Maps the mysql client's errors back to statements and returns a list
298
+ of (error code, message), with both entries None on success."""
299
+ results: list[tuple[str | None, str | None]] = [(None, None)] * count
300
+ for raw in stderr.splitlines():
301
+ m = ERROR.match(raw)
302
+ if not m:
303
+ continue
304
+ at = int(m.group(2))
305
+ i = bisect.bisect_right(starts, at) - 1
306
+ if i < 0:
307
+ raise RuntimeError(f"replay setup failed: {raw}")
308
+ if results[i][0] is None:
309
+ results[i] = (m.group(1), m.group(3))
310
+ return results
311
+
312
+
313
+ def _entries(log):
314
+ """Yields (connection ID, command, argument) for each general log entry."""
315
+ entry = None
316
+ for raw in log:
317
+ m = ENTRY.match(raw)
318
+ if m:
319
+ if entry:
320
+ yield _decode(entry)
321
+ entry = [m.group(1), m.group(2), m.group(3)]
322
+ elif entry:
323
+ entry[2] += raw
324
+ if entry:
325
+ yield _decode(entry)
326
+
327
+
328
+ def _decode(entry):
329
+ conn, command, arg = entry
330
+ # Statements can carry binary values that aren't valid UTF-8.
331
+ # surrogateescape keeps those bytes intact through to replay.
332
+ return int(conn), command.decode().strip(), arg.decode("utf-8", "surrogateescape").removesuffix("\n")
333
+
334
+
335
+ def _lenient_truncate(name):
336
+ """Rewrites a TRUNCATE to call _migcheck_truncate, which replay defines."""
337
+ table = name.split(".")[-1].strip("`")
338
+ return "CALL _migcheck_truncate('" + table.replace("\\", "\\\\").replace("'", "''") + "')"