diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d91317..ee430ce 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## Unreleased +### Features + +- Add JSON output for batch queries with the `--json` option. + ### Bug Fixes - Expand `~` in configured log file paths before opening the log. diff --git a/litecli/main.py b/litecli/main.py index fa732c3..2a9d4c1 100644 --- a/litecli/main.py +++ b/litecli/main.py @@ -16,7 +16,7 @@ import click import sqlparse -from cli_helpers.tabular_output import TabularOutputFormatter, preprocessors +from cli_helpers.tabular_output import TabularOutputFormatter, json_output_adapter, preprocessors from prompt_toolkit.auto_suggest import AutoSuggestFromHistory from prompt_toolkit.completion import Completion, DynamicCompleter from prompt_toolkit.document import Document @@ -67,6 +67,18 @@ def _load_sqlite3() -> Any: PACKAGE_ROOT = os.path.abspath(os.path.dirname(__file__)) +def _json_output_adapter(data: Iterable[Iterable[Any]], headers: Iterable[Any], **kwargs: Any) -> str: + rows = json_output_adapter.adapter(data, headers, **kwargs) + return "[{}]".format(",".join(rows)) + + +TabularOutputFormatter.register_new_formatter( + "json", + _json_output_adapter, + json_output_adapter.preprocessors, +) + + class LiteCli(object): default_prompt = "\\d> " max_len_prompt = 45 @@ -931,6 +943,7 @@ def get_last_query(self) -> str | None: ) @click.option("-t", "--table", is_flag=True, help="Display batch output in table format.") @click.option("--csv", is_flag=True, help="Display batch output in CSV format.") +@click.option("--json", "json_output", is_flag=True, help="Display batch output in JSON format.") @click.option("--warn/--no-warn", default=None, help="Warn before running a destructive query.") @click.option("-e", "--execute", type=str, help="Execute command and quit.") @click.argument("database", default="", nargs=1) @@ -942,6 +955,7 @@ def cli( auto_vertical_output: bool, table: bool, csv: bool, + json_output: bool, warn: bool | None, execute: str | None, liteclirc: str, @@ -971,7 +985,9 @@ def cli( # --execute argument if execute: try: - if csv: + if json_output: + litecli.formatter.format_name = "json" + elif csv: litecli.formatter.format_name = "csv" elif not table: litecli.formatter.format_name = "tsv" @@ -998,7 +1014,9 @@ def cli( try: new_line = True - if csv: + if json_output: + litecli.formatter.format_name = "json" + elif csv: litecli.formatter.format_name = "csv" elif not table: litecli.formatter.format_name = "tsv" diff --git a/tests/test_main.py b/tests/test_main.py index 21ad5b6..e0db383 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -76,6 +76,14 @@ def test_execute_arg_with_csv(executor): assert expected in "".join(result.output) +def test_execute_arg_with_json(): + runner = CliRunner() + result = runner.invoke(cli, args=CLI_ARGS[:-1] + [":memory:", "-e", "select 'abc' as a;", "--json"]) + + assert result.exit_code == 0 + assert '[{"a":"abc"}]' in result.output + + @dbtest def test_batch_mode(executor): run(executor, """create table test(a text)""") @@ -134,6 +142,14 @@ def test_batch_mode_csv(executor): assert expected in "".join(result.output) +def test_batch_mode_json(): + runner = CliRunner() + result = runner.invoke(cli, args=CLI_ARGS[:-1] + [":memory:", "--json"], input="select 'abc' as a;") + + assert result.exit_code == 0 + assert '[{"a":"abc"}]' in result.output + + def test_help_strings_end_with_periods(): """Make sure click options have help text that end with a period.""" for param in cli.params: