diff --git a/AUTHORS b/AUTHORS index db49b9eb4..691396045 100644 --- a/AUTHORS +++ b/AUTHORS @@ -145,6 +145,7 @@ Contributors: * fbdb * Charbel Jacquin (charbeljc) * Jeronimo Garcia (bechampion) + * Diego * Devadathan M B (devadathanmb) * Charalampos Stratakis * Laszlo Bimba (bimlas) diff --git a/changelog.rst b/changelog.rst index fbdbc4917..ee5d1fcf5 100644 --- a/changelog.rst +++ b/changelog.rst @@ -3,6 +3,11 @@ Upcoming Features: --------- +* Add support for `single-command` to run a SQL command and exit. + * Command line option `-c` or `--command`. + * You can specify multiple times. + * Runs one statement at a time, like `-f`, and can be combined with `-f`: + both run, the same way psql does. * Add support for forcing destructive commands without confirmation. * Command line option `-y` or `--yes`. * Skips the destructive command confirmation prompt when enabled. diff --git a/pgcli/main.py b/pgcli/main.py index 8c172b85d..3ff820ce2 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -1057,6 +1057,29 @@ def _check_ongoing_transaction_and_allow_quitting(self): def run_cli(self): logger = self.logger + # Handle command mode (-c flag) - similar to psql behavior + # Multiple -c options are executed sequentially + if hasattr(self, 'commands') and self.commands: + try: + for command in self.commands: + logger.debug("Running command: %s", command) + # Statement by statement, like psql -c: \watch only repeats + # its own statement, not the whole -c block. + if not self._execute_statements(command): + break + except PgCliQuitError: + # Normal exit from quit command + sys.exit(0) + except Exception as e: + logger.error("Error executing command: %s", e) + logger.error("traceback: %r", traceback.format_exc()) + click.secho(str(e), err=True, fg="red") + sys.exit(1) + # psql runs both -c and -f when they are given together, so only + # exit here when there is no file left to run. + if not (hasattr(self, 'input_files') and self.input_files): + sys.exit(0) + # Handle file mode (-f flag) - similar to psql behavior # Multiple -f options are executed sequentially if hasattr(self, 'input_files') and self.input_files: @@ -1503,8 +1526,13 @@ def is_too_tall(self, lines): return len(lines) >= (self.prompt_app.output.get_size().rows - 4) def echo_via_pager(self, text, color=None): - # Disable pager for -f/--file mode and \watch command - if self.pgspecial.pager_config == PAGER_OFF or self.watch_command or (hasattr(self, 'input_files') and self.input_files): + # Disable pager for -c/--command and -f/--file modes and \watch command + if ( + self.pgspecial.pager_config == PAGER_OFF + or self.watch_command + or (hasattr(self, 'commands') and self.commands) + or (hasattr(self, 'input_files') and self.input_files) + ): click.echo(text, color=color) elif self.pgspecial.pager_config == PAGER_LONG_OUTPUT and self.table_format != "csv": lines = text.split("\n") @@ -1667,6 +1695,13 @@ def echo_via_pager(self, text, color=None): type=str, help="SQL statement to execute after connecting.", ) +@click.option( + "-c", + "--command", + "commands", + multiple=True, + help="run command (SQL or internal) and exit. Multiple -c options are allowed.", +) @click.option( "-y", "--yes", @@ -1712,6 +1747,7 @@ def cli( ssh_tunnel: str, init_command: str, log_file: str, + commands: tuple, force_destructive: bool, input_files: tuple, connect_timeout: int | None, @@ -1783,6 +1819,8 @@ def cli( connect_timeout=connect_timeout, ) + # Store commands for -c option (can be multiple) + pgcli.commands = commands if commands else None # Store file paths for -f option (can be multiple) pgcli.input_files = input_files if input_files else None diff --git a/tests/features/command_option.feature b/tests/features/command_option.feature new file mode 100644 index 000000000..38f5fef51 --- /dev/null +++ b/tests/features/command_option.feature @@ -0,0 +1,38 @@ +Feature: run the cli with -c/--command option, + execute a single command, + and exit + + Scenario: run pgcli with -c and a SQL query + When we run pgcli with -c "SELECT 1 as test_column" + then we see the query result + and pgcli exits successfully + + Scenario: run pgcli with --command and a SQL query + When we run pgcli with --command "SELECT 'hello' as greeting" + then we see the query result + and pgcli exits successfully + + Scenario: run pgcli with -c and a special command + When we run pgcli with -c "\dt" + then we see the command output + and pgcli exits successfully + + Scenario: run pgcli with -c and an invalid query + When we run pgcli with -c "SELECT invalid_column FROM nonexistent_table" + then we see an error message + and pgcli exits successfully + + Scenario: run pgcli with -c and multiple statements + When we run pgcli with -c "SELECT 1; SELECT 2" + then we see both query results + and pgcli exits successfully + + Scenario: run pgcli with multiple -c options + When we run pgcli with multiple -c options + then we see all command outputs + and pgcli exits successfully + + Scenario: run pgcli with mixed -c and --command options + When we run pgcli with mixed -c and --command + then we see all command outputs + and pgcli exits successfully diff --git a/tests/features/steps/command_option.py b/tests/features/steps/command_option.py new file mode 100644 index 000000000..b783dc7a7 --- /dev/null +++ b/tests/features/steps/command_option.py @@ -0,0 +1,144 @@ +"""Steps for the -c/--command option. + +The steps shared with the -f/--file scenarios (query result, command output, +error message, exit status) are defined once in file_option.py; behave's step +registry is global, so redefining them here would raise AmbiguousStep. +""" + +import subprocess +from behave import when, then + + +@when('we run pgcli with -c "{command}"') +def step_run_pgcli_with_c(context, command): + """Run pgcli with -c flag and a command.""" + cmd = [ + "pgcli", + "-h", + context.conf["host"], + "-p", + str(context.conf["port"]), + "-U", + context.conf["user"], + "-d", + context.conf["dbname"], + "-c", + command, + ] + try: + context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=5) + context.exit_code = 0 + except subprocess.CalledProcessError as e: + context.cmd_output = e.output + context.exit_code = e.returncode + except subprocess.TimeoutExpired: + context.cmd_output = b"Command timed out" + context.exit_code = -1 + + +@when('we run pgcli with --command "{command}"') +def step_run_pgcli_with_command(context, command): + """Run pgcli with --command flag and a command.""" + cmd = [ + "pgcli", + "-h", + context.conf["host"], + "-p", + str(context.conf["port"]), + "-U", + context.conf["user"], + "-d", + context.conf["dbname"], + "--command", + command, + ] + try: + context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=5) + context.exit_code = 0 + except subprocess.CalledProcessError as e: + context.cmd_output = e.output + context.exit_code = e.returncode + except subprocess.TimeoutExpired: + context.cmd_output = b"Command timed out" + context.exit_code = -1 + + +@then("pgcli exits with error") +def step_pgcli_exits_with_error(context): + """Verify that pgcli exited with a non-zero code.""" + assert context.exit_code != 0, f"Expected non-zero exit code, but got: {context.exit_code}" + # Clean up + context.cmd_output = None + context.exit_code = None + + +@when("we run pgcli with multiple -c options") +def step_run_pgcli_with_multiple_c(context): + """Run pgcli with multiple -c flags.""" + cmd = [ + "pgcli", + "-h", + context.conf["host"], + "-p", + str(context.conf["port"]), + "-U", + context.conf["user"], + "-d", + context.conf["dbname"], + "-c", + "SELECT 'first' as result", + "-c", + "SELECT 'second' as result", + "-c", + "SELECT 'third' as result", + ] + try: + context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=10) + context.exit_code = 0 + except subprocess.CalledProcessError as e: + context.cmd_output = e.output + context.exit_code = e.returncode + except subprocess.TimeoutExpired: + context.cmd_output = b"Command timed out" + context.exit_code = -1 + + +@when("we run pgcli with mixed -c and --command") +def step_run_pgcli_with_mixed_options(context): + """Run pgcli with mixed -c and --command flags.""" + cmd = [ + "pgcli", + "-h", + context.conf["host"], + "-p", + str(context.conf["port"]), + "-U", + context.conf["user"], + "-d", + context.conf["dbname"], + "-c", + "SELECT 'from_c' as source", + "--command", + "SELECT 'from_command' as source", + ] + try: + context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=10) + context.exit_code = 0 + except subprocess.CalledProcessError as e: + context.cmd_output = e.output + context.exit_code = e.returncode + except subprocess.TimeoutExpired: + context.cmd_output = b"Command timed out" + context.exit_code = -1 + + +@then("we see all command outputs") +def step_see_all_command_outputs(context): + """Verify that all command outputs are present.""" + output = context.cmd_output.decode('utf-8') + # Should contain output from all commands + assert "first" in output or "from_c" in output, f"Expected 'first' or 'from_c' in output, but got: {output}" + assert "second" in output or "from_command" in output, f"Expected 'second' or 'from_command' in output, but got: {output}" + # For the 3-command test, also check for third + if "third" in output or "result" in output: + assert "third" in output, f"Expected 'third' in output for 3-command test, but got: {output}" diff --git a/tests/test_main.py b/tests/test_main.py index c8a28b419..f18319312 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -563,6 +563,31 @@ def test_execute_statements_does_not_split_inside_literals(executor): assert "a;b" in outputs[0] +def test_command_and_file_both_run(tmpdir): + """psql runs both -c and -f when given together; -c must not exit first.""" + sql_file = tmpdir.join("script.sql") + sql_file.write("select 2;") + cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile"))) + cli.commands = ["select 1;"] + cli.input_files = [str(sql_file)] + with mock.patch.object(cli, "_execute_statements", return_value=True) as mock_exec: + with pytest.raises(SystemExit) as e: + cli.run_cli() + assert e.value.code == 0 + assert [c[0][0] for c in mock_exec.call_args_list] == ["select 1;", "select 2;"] + + +def test_command_mode_alone_still_exits(tmpdir): + """With no -f, the -c block still exits on its own.""" + cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile"))) + cli.commands = ["select 1;"] + with mock.patch.object(cli, "_execute_statements", return_value=True) as mock_exec: + with pytest.raises(SystemExit) as e: + cli.run_cli() + assert e.value.code == 0 + mock_exec.assert_called_once_with("select 1;") + + def test_file_mode_runs_statements(tmpdir): """-f wiring: the file content goes through _execute_statements.""" sql_file = tmpdir.join("script.sql")