diff --git a/pipekit/jrunner.py b/pipekit/jrunner.py index b72d91c..dc36229 100644 --- a/pipekit/jrunner.py +++ b/pipekit/jrunner.py @@ -14,9 +14,11 @@ argv or in the database. from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Iterator +import contextlib import csv import io +import json import os import re import shutil @@ -79,6 +81,36 @@ def jrunner_path() -> Path: return get_config().jrunner_path +def _write_passfile(entries: dict[str, dict]) -> str: + """Write a jrunner alias file for `entries` and return its path. + + ``entries`` maps alias -> connection dict. Passwords resolve here and reach + jrunner through this file rather than through ``-scp``/``-dcp``, because + anything on argv is readable by any local user via ``ps`` for the lifetime + of the call. ``mkstemp`` creates the file 0600. + + The caller is responsible for deleting it — every call site already has a + ``finally`` that removes the temporary SQL file, so cleanup rides along + there rather than forcing the whole body into another ``with``. + """ + fd, path = tempfile.mkstemp(suffix=".jrunnerpass") + with os.fdopen(fd, "w") as f: + for alias, conn in entries.items(): + f.write(f"[{alias}]\n") + f.write(f"url={conn['jdbc_url']}\n") + f.write(f"user={conn.get('username') or ''}\n") + f.write(f"pass={resolve_password(conn.get('password'))}\n") + return path + + +def _unlink_quietly(*paths: str) -> None: + for p in paths: + try: + os.unlink(p) + except (FileNotFoundError, TypeError): + pass + + def version() -> tuple[bool, str]: """Return (ok, message) for use by pipekit doctor.""" path = jrunner_path() @@ -106,15 +138,16 @@ def query( ) -> QueryResult: """Run `sql` in jrunner query mode and parse CSV output.""" path = jrunner_path() - pw = resolve_password(password) + conn = {"jdbc_url": jdbc_url, "username": username, "password": password} with tempfile.NamedTemporaryFile("w", suffix=".sql", delete=False) as f: f.write(sql) sql_path = f.name + pass_path = _write_passfile({"src": conn}) try: argv = [str(path), - "-scu", jdbc_url, - "-scn", username or "", - "-scp", pw, + "-sc", "src", + "--passfile", pass_path, + "--strict", "-sq", sql_path, "-f", "csv"] if trim: @@ -122,7 +155,7 @@ def query( r = subprocess.run(argv, capture_output=True, text=True, timeout=timeout, env=_subprocess_env()) finally: - os.unlink(sql_path) + _unlink_quietly(sql_path, pass_path) if r.returncode != 0: raise JrunnerError(r.stderr.strip() or r.stdout.strip(), @@ -161,14 +194,13 @@ def migrate( with tempfile.NamedTemporaryFile("w", suffix=".sql", delete=False) as f: f.write(sql) sql_path = f.name + pass_path = _write_passfile({"src": source_conn, "dst": dest_conn}) try: argv = [str(path), - "-scu", source_conn["jdbc_url"], - "-scn", source_conn.get("username") or "", - "-scp", resolve_password(source_conn.get("password")), - "-dcu", dest_conn["jdbc_url"], - "-dcn", dest_conn.get("username") or "", - "-dcp", resolve_password(dest_conn.get("password")), + "-sc", "src", + "-dc", "dst", + "--passfile", pass_path, + "--strict", "-sq", sql_path, "-dt", dest_table] if trim: @@ -215,7 +247,7 @@ def migrate( stdout = "".join(stdout_lines) stderr = "".join(stderr_buf) finally: - os.unlink(sql_path) + _unlink_quietly(sql_path, pass_path) if proc.returncode != 0: raise JrunnerError(stderr.strip() or stdout.strip(), @@ -224,10 +256,14 @@ def migrate( if silent: raise JrunnerError(silent, stdout=stdout, stderr=stderr) - return MigrateResult( - row_count=_parse_row_count(stdout + "\n" + stderr), - stdout=stdout, stderr=stderr, - ) + summary = _parse_summary(stderr) + row_count = (summary or {}).get("rows") + if row_count is None: + # Older jrunner without @summary — fall back to reading the count out of + # its prose. Removable once the pinned jar is the only one in play. + row_count = _parse_row_count(stdout + "\n" + stderr) + + return MigrateResult(row_count=row_count, stdout=stdout, stderr=stderr) def run_dest_sql(conn: dict, sql: str, *, timeout: int = 600) -> QueryResult: @@ -238,6 +274,27 @@ def run_dest_sql(conn: dict, sql: str, *, timeout: int = 600) -> QueryResult: sql, timeout=timeout, trim=False) +_SUMMARY_PREFIX = "@summary " + + +def _parse_summary(stderr: str) -> dict | None: + """Parse jrunner's ``@summary {...}`` line out of stderr. + + jrunner emits exactly one, last, carrying ``status`` and (where known) + ``rows`` and ``ms``. It lives on stderr because migration-mode stdout is + human-readable progress and query-mode stdout is CSV data. Returns None if + absent, which is how an older jrunner presents. + """ + for line in reversed((stderr or "").splitlines()): + line = line.strip() + if line.startswith(_SUMMARY_PREFIX): + try: + return json.loads(line[len(_SUMMARY_PREFIX):]) + except ValueError: + return None + return None + + _ROW_COUNT_PATTERNS = ( re.compile(r"(\d+)\s+rows?\s+(?:inserted|transferred|migrated|written)", re.I), re.compile(r"inserted\s+(\d+)\s+rows?", re.I),