Skip --checkpoint file statements with --resume

changelog.md

@@ -5,6 +5,7 @@ Features
 ---------
 * Remove undocumented `%mycli` Jupyter magic.
 * Add `--quiet` option, and let `--verbose` be given multiple times.
+* Add `--resume` to replay `--checkpoint` files with `--batch`.
 
 
 Bug Fixes

mycli/main.py

@@ -1354,7 +1354,12 @@ class CliArgs:
     )
     checkpoint: TextIOWrapper | None = clickdc.option(
         type=click.File(mode='a', encoding='utf-8'),
-        help='In batch or --execute mode, log successful queries to a file.',
+        help='In batch or --execute mode, log successful queries to a file, and skipped with --resume.',
+    )
+    resume: bool = clickdc.option(
+        '--resume',
+        is_flag=True,
+        help='In batch mode, resume after replaying statements in the --checkpoint file.',
     )
     defaults_group_suffix: str | None = clickdc.option(
         type=str,
@@ -1522,6 +1527,14 @@ def get_password_from_file(password_file: str | None) -> str | None:
     if cli_args.password is None and os.environ.get("MYSQL_PWD") is not None:
         cli_args.password = os.environ.get("MYSQL_PWD")
 
+    if cli_args.resume and not cli_args.checkpoint:
+        click.secho('Error: --resume requires a --checkpoint file.', err=True, fg='red')
+        sys.exit(1)
+
+    if cli_args.resume and not cli_args.batch:
+        click.secho('Error: --resume requires a --batch file.', err=True, fg='red')
+        sys.exit(1)
+
     cli_verbosity = 0
     if cli_args.verbose and cli_args.quiet:
         click.secho('Error: --verbose and --quiet are incompatible.', err=True, fg='red')

mycli/main_modes/batch.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from io import TextIOWrapper
 import os
 import sys
 import time
@@ -19,6 +20,53 @@
     from mycli.main import CliArgs, MyCli
 
 
+class CheckpointReplayError(Exception):
+    pass
+
+
+def replay_checkpoint_file(
+    batch_path: str,
+    checkpoint: TextIOWrapper | None,
+    resume: bool,
+) -> int:
+    if not resume:
+        return 0
+
+    if checkpoint is None:
+        return 0
+
+    if batch_path == '-':
+        raise CheckpointReplayError('--resume is incompatible with reading from the standard input.')
+
+    checkpoint_name = checkpoint.name
+    checkpoint.flush()
+    completed_count = 0
+    try:
+        with click.open_file(batch_path) as batch_h, click.open_file(checkpoint_name, mode='r', encoding='utf-8') as checkpoint_h:
+            try:
+                batch_gen = statements_from_filehandle(batch_h)
+            except ValueError as e:
+                raise CheckpointReplayError(f'Error reading --batch file: {batch_path}: {e}') from None
+            for checkpoint_statement, _checkpoint_counter in statements_from_filehandle(checkpoint_h):
+                try:
+                    batch_statement, _batch_counter = next(batch_gen)
+                except StopIteration:
+                    raise CheckpointReplayError('Checkpoint script longer than batch script.') from None
+                except ValueError as e:
+                    raise CheckpointReplayError(f'Error reading --batch file: {batch_path}: {e}') from None
+                if checkpoint_statement != batch_statement:
+                    raise CheckpointReplayError(f'Statement mismatch: {checkpoint_statement}.')
+                completed_count += 1
+    except ValueError as e:
+        raise CheckpointReplayError(f'Error reading --checkpoint file: {checkpoint.name}: {e}') from None
+    except FileNotFoundError as e:
+        raise CheckpointReplayError(f'FileNotFoundError: {e}') from None
+    except OSError as e:
+        raise CheckpointReplayError(f'OSError: {e}') from None
+
+    return completed_count
+
+
 def dispatch_batch_statements(
     mycli: 'MyCli',
     cli_args: 'CliArgs',
@@ -70,6 +118,7 @@ def main_batch_with_progress_bar(mycli: 'MyCli', cli_args: 'CliArgs') -> int:
         click.secho('--progress is only compatible with a plain file.', err=True, fg='red')
         return 1
     try:
+        completed_statement_count = replay_checkpoint_file(cli_args.batch, cli_args.checkpoint, cli_args.resume)
         batch_count_h = click.open_file(cli_args.batch)
         for _statement, _counter in statements_from_filehandle(batch_count_h):
             goal_statements += 1
@@ -82,6 +131,10 @@ def main_batch_with_progress_bar(mycli: 'MyCli', cli_args: 'CliArgs') -> int:
     except ValueError as e:
         click.secho(f'Error reading --batch file: {cli_args.batch}: {e}', err=True, fg='red')
         return 1
+    except CheckpointReplayError as e:
+        name = cli_args.checkpoint.name if cli_args.checkpoint else 'None'
+        click.secho(f'Error replaying --checkpoint file: {name}: {e}', err=True, fg='red')
+        return 1
     try:
         if goal_statements:
             pb_style = prompt_toolkit.styles.Style.from_dict({'bar-a': 'reverse'})
@@ -98,6 +151,8 @@ def main_batch_with_progress_bar(mycli: 'MyCli', cli_args: 'CliArgs') -> int:
             with ProgressBar(style=pb_style, formatters=custom_formatters, output=err_output) as pb:
                 for _pb_counter in pb(range(goal_statements)):
                     statement, statement_counter = next(batch_gen)
+                    if statement_counter < completed_statement_count:
+                        continue
                     dispatch_batch_statements(mycli, cli_args, statement, statement_counter)
     except (ValueError, StopIteration, IOError, OSError, pymysql.err.Error) as e:
         click.secho(str(e), err=True, fg='red')
@@ -113,12 +168,19 @@ def main_batch_without_progress_bar(mycli: 'MyCli', cli_args: 'CliArgs') -> int:
     if not sys.stdin.isatty() and cli_args.batch != '-':
         click.secho('Ignoring STDIN since --batch was also given.', err=True, fg='red')
     try:
+        completed_statement_count = replay_checkpoint_file(cli_args.batch, cli_args.checkpoint, cli_args.resume)
         batch_h = click.open_file(cli_args.batch)
     except (OSError, FileNotFoundError):
         click.secho(f'Failed to open --batch file: {cli_args.batch}', err=True, fg='red')
         return 1
+    except CheckpointReplayError as e:
+        name = cli_args.checkpoint.name if cli_args.checkpoint else 'None'
+        click.secho(f'Error replaying --checkpoint file: {name}: {e}', err=True, fg='red')
+        return 1
     try:
         for statement, counter in statements_from_filehandle(batch_h):
+            if counter < completed_statement_count:
+                continue
             dispatch_batch_statements(mycli, cli_args, statement, counter)
     except (ValueError, StopIteration, IOError, OSError, pymysql.err.Error) as e:
         click.secho(str(e), err=True, fg='red')

test/pytests/test_main.py

@@ -2156,6 +2156,15 @@ def test_quiet_sets_negative_cli_verbosity(monkeypatch: pytest.MonkeyPatch) -> N
     assert dummy.init_kwargs['cli_verbosity'] == -1
 
 
+def test_resume_requires_checkpoint() -> None:
+    runner = CliRunner()
+
+    result = runner.invoke(click_entrypoint, args=['--batch', os.devnull, '--resume'])
+
+    assert result.exit_code == 1
+    assert 'Error:' in result.output
+
+
 def test_execute_arg_supersedes_batch_file(monkeypatch):
     mycli_main, mycli_main_batch, MockMyCli = noninteractive_mock_mycli(monkeypatch)
     runner = CliRunner()