Skip to content

Commit 924e7d4

Browse files
DiegoDAFclaudediego-feito-stori
authored
Add support for -c/--command option to execute commands and exit (#1542)
* Add support for multiple -c/--command parameters This commit adds support for the -c/--command option to pgcli, similar to psql's behavior. Users can now execute SQL commands or internal commands from the command line and exit immediately after execution. Features: - Single command execution: pgcli -c "SELECT 1" - Multiple commands: pgcli -c "SELECT 1" -c "SELECT 2" - Long form: pgcli --command "SELECT 1" - Mixed usage: pgcli -c "cmd1" --command "cmd2" - Commands are executed sequentially - Pager is automatically disabled in command mode - Proper error handling and exit codes Tests included for all scenarios. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Assert on the \dt output in the command-output step Address review feedback: the step decoded cmd_output into `output` but never used it (flagged by CodeQL as an unused local variable). It now asserts on the \dt column headers, which are rendered whether or not any tables exist, so the step verifies the special command actually produced its listing. Also drop the unused `as e` bindings in the TimeoutExpired handlers (ruff F841) and apply ruff format. * Reuse the shared behave steps instead of redefining them behave's step registry is global, so now that -f landed and brought tests/features/steps/file_option.py with it, the five steps this file also defined ("we see the query result", "we see both query results", "we see the command output", "we see an error message", "pgcli exits successfully") raised AmbiguousStep and aborted the whole integration run before any scenario executed. Four of the five were byte-identical anyway; the fifth only differed in a column name it looks for, and its assertion is an any() that already matches on "SELECT". Dropping the duplicates here and reusing the ones from file_option.py keeps every -c scenario passing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: DiegoDAF <DiegoDAF@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Diego - pgDBA <151667121+diego-feito-stori@users.noreply.github.com>
1 parent 9bd366a commit 924e7d4

6 files changed

Lines changed: 253 additions & 2 deletions

File tree

‎AUTHORS‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,7 @@ Contributors:
145145
* fbdb
146146
* Charbel Jacquin (charbeljc)
147147
* Jeronimo Garcia (bechampion)
148+
* Diego
148149
* Devadathan M B (devadathanmb)
149150
* Charalampos Stratakis
150151
* Laszlo Bimba (bimlas)

‎changelog.rst‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,11 @@ Upcoming
33

44
Features:
55
---------
6+
* Add support for `single-command` to run a SQL command and exit.
7+
* Command line option `-c` or `--command`.
8+
* You can specify multiple times.
9+
* Runs one statement at a time, like `-f`, and can be combined with `-f`:
10+
both run, the same way psql does.
611
* Add support for forcing destructive commands without confirmation.
712
* Command line option `-y` or `--yes`.
813
* Skips the destructive command confirmation prompt when enabled.

‎pgcli/main.py‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1057,6 +1057,29 @@ def _check_ongoing_transaction_and_allow_quitting(self):
10571057
def run_cli(self):
10581058
logger = self.logger
10591059

1060+
# Handle command mode (-c flag) - similar to psql behavior
1061+
# Multiple -c options are executed sequentially
1062+
if hasattr(self, 'commands') and self.commands:
1063+
try:
1064+
for command in self.commands:
1065+
logger.debug("Running command: %s", command)
1066+
# Statement by statement, like psql -c: \watch only repeats
1067+
# its own statement, not the whole -c block.
1068+
if not self._execute_statements(command):
1069+
break
1070+
except PgCliQuitError:
1071+
# Normal exit from quit command
1072+
sys.exit(0)
1073+
except Exception as e:
1074+
logger.error("Error executing command: %s", e)
1075+
logger.error("traceback: %r", traceback.format_exc())
1076+
click.secho(str(e), err=True, fg="red")
1077+
sys.exit(1)
1078+
# psql runs both -c and -f when they are given together, so only
1079+
# exit here when there is no file left to run.
1080+
if not (hasattr(self, 'input_files') and self.input_files):
1081+
sys.exit(0)
1082+
10601083
# Handle file mode (-f flag) - similar to psql behavior
10611084
# Multiple -f options are executed sequentially
10621085
if hasattr(self, 'input_files') and self.input_files:
@@ -1503,8 +1526,13 @@ def is_too_tall(self, lines):
15031526
return len(lines) >= (self.prompt_app.output.get_size().rows - 4)
15041527

15051528
def echo_via_pager(self, text, color=None):
1506-
# Disable pager for -f/--file mode and \watch command
1507-
if self.pgspecial.pager_config == PAGER_OFF or self.watch_command or (hasattr(self, 'input_files') and self.input_files):
1529+
# Disable pager for -c/--command and -f/--file modes and \watch command
1530+
if (
1531+
self.pgspecial.pager_config == PAGER_OFF
1532+
or self.watch_command
1533+
or (hasattr(self, 'commands') and self.commands)
1534+
or (hasattr(self, 'input_files') and self.input_files)
1535+
):
15081536
click.echo(text, color=color)
15091537
elif self.pgspecial.pager_config == PAGER_LONG_OUTPUT and self.table_format != "csv":
15101538
lines = text.split("\n")
@@ -1667,6 +1695,13 @@ def echo_via_pager(self, text, color=None):
16671695
type=str,
16681696
help="SQL statement to execute after connecting.",
16691697
)
1698+
@click.option(
1699+
"-c",
1700+
"--command",
1701+
"commands",
1702+
multiple=True,
1703+
help="run command (SQL or internal) and exit. Multiple -c options are allowed.",
1704+
)
16701705
@click.option(
16711706
"-y",
16721707
"--yes",
@@ -1712,6 +1747,7 @@ def cli(
17121747
ssh_tunnel: str,
17131748
init_command: str,
17141749
log_file: str,
1750+
commands: tuple,
17151751
force_destructive: bool,
17161752
input_files: tuple,
17171753
connect_timeout: int | None,
@@ -1783,6 +1819,8 @@ def cli(
17831819
connect_timeout=connect_timeout,
17841820
)
17851821

1822+
# Store commands for -c option (can be multiple)
1823+
pgcli.commands = commands if commands else None
17861824
# Store file paths for -f option (can be multiple)
17871825
pgcli.input_files = input_files if input_files else None
17881826

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
Feature: run the cli with -c/--command option,
2+
execute a single command,
3+
and exit
4+
5+
Scenario: run pgcli with -c and a SQL query
6+
When we run pgcli with -c "SELECT 1 as test_column"
7+
then we see the query result
8+
and pgcli exits successfully
9+
10+
Scenario: run pgcli with --command and a SQL query
11+
When we run pgcli with --command "SELECT 'hello' as greeting"
12+
then we see the query result
13+
and pgcli exits successfully
14+
15+
Scenario: run pgcli with -c and a special command
16+
When we run pgcli with -c "\dt"
17+
then we see the command output
18+
and pgcli exits successfully
19+
20+
Scenario: run pgcli with -c and an invalid query
21+
When we run pgcli with -c "SELECT invalid_column FROM nonexistent_table"
22+
then we see an error message
23+
and pgcli exits successfully
24+
25+
Scenario: run pgcli with -c and multiple statements
26+
When we run pgcli with -c "SELECT 1; SELECT 2"
27+
then we see both query results
28+
and pgcli exits successfully
29+
30+
Scenario: run pgcli with multiple -c options
31+
When we run pgcli with multiple -c options
32+
then we see all command outputs
33+
and pgcli exits successfully
34+
35+
Scenario: run pgcli with mixed -c and --command options
36+
When we run pgcli with mixed -c and --command
37+
then we see all command outputs
38+
and pgcli exits successfully
Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
"""Steps for the -c/--command option.
2+
3+
The steps shared with the -f/--file scenarios (query result, command output,
4+
error message, exit status) are defined once in file_option.py; behave's step
5+
registry is global, so redefining them here would raise AmbiguousStep.
6+
"""
7+
8+
import subprocess
9+
from behave import when, then
10+
11+
12+
@when('we run pgcli with -c "{command}"')
13+
def step_run_pgcli_with_c(context, command):
14+
"""Run pgcli with -c flag and a command."""
15+
cmd = [
16+
"pgcli",
17+
"-h",
18+
context.conf["host"],
19+
"-p",
20+
str(context.conf["port"]),
21+
"-U",
22+
context.conf["user"],
23+
"-d",
24+
context.conf["dbname"],
25+
"-c",
26+
command,
27+
]
28+
try:
29+
context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=5)
30+
context.exit_code = 0
31+
except subprocess.CalledProcessError as e:
32+
context.cmd_output = e.output
33+
context.exit_code = e.returncode
34+
except subprocess.TimeoutExpired:
35+
context.cmd_output = b"Command timed out"
36+
context.exit_code = -1
37+
38+
39+
@when('we run pgcli with --command "{command}"')
40+
def step_run_pgcli_with_command(context, command):
41+
"""Run pgcli with --command flag and a command."""
42+
cmd = [
43+
"pgcli",
44+
"-h",
45+
context.conf["host"],
46+
"-p",
47+
str(context.conf["port"]),
48+
"-U",
49+
context.conf["user"],
50+
"-d",
51+
context.conf["dbname"],
52+
"--command",
53+
command,
54+
]
55+
try:
56+
context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=5)
57+
context.exit_code = 0
58+
except subprocess.CalledProcessError as e:
59+
context.cmd_output = e.output
60+
context.exit_code = e.returncode
61+
except subprocess.TimeoutExpired:
62+
context.cmd_output = b"Command timed out"
63+
context.exit_code = -1
64+
65+
66+
@then("pgcli exits with error")
67+
def step_pgcli_exits_with_error(context):
68+
"""Verify that pgcli exited with a non-zero code."""
69+
assert context.exit_code != 0, f"Expected non-zero exit code, but got: {context.exit_code}"
70+
# Clean up
71+
context.cmd_output = None
72+
context.exit_code = None
73+
74+
75+
@when("we run pgcli with multiple -c options")
76+
def step_run_pgcli_with_multiple_c(context):
77+
"""Run pgcli with multiple -c flags."""
78+
cmd = [
79+
"pgcli",
80+
"-h",
81+
context.conf["host"],
82+
"-p",
83+
str(context.conf["port"]),
84+
"-U",
85+
context.conf["user"],
86+
"-d",
87+
context.conf["dbname"],
88+
"-c",
89+
"SELECT 'first' as result",
90+
"-c",
91+
"SELECT 'second' as result",
92+
"-c",
93+
"SELECT 'third' as result",
94+
]
95+
try:
96+
context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=10)
97+
context.exit_code = 0
98+
except subprocess.CalledProcessError as e:
99+
context.cmd_output = e.output
100+
context.exit_code = e.returncode
101+
except subprocess.TimeoutExpired:
102+
context.cmd_output = b"Command timed out"
103+
context.exit_code = -1
104+
105+
106+
@when("we run pgcli with mixed -c and --command")
107+
def step_run_pgcli_with_mixed_options(context):
108+
"""Run pgcli with mixed -c and --command flags."""
109+
cmd = [
110+
"pgcli",
111+
"-h",
112+
context.conf["host"],
113+
"-p",
114+
str(context.conf["port"]),
115+
"-U",
116+
context.conf["user"],
117+
"-d",
118+
context.conf["dbname"],
119+
"-c",
120+
"SELECT 'from_c' as source",
121+
"--command",
122+
"SELECT 'from_command' as source",
123+
]
124+
try:
125+
context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root, stderr=subprocess.STDOUT, timeout=10)
126+
context.exit_code = 0
127+
except subprocess.CalledProcessError as e:
128+
context.cmd_output = e.output
129+
context.exit_code = e.returncode
130+
except subprocess.TimeoutExpired:
131+
context.cmd_output = b"Command timed out"
132+
context.exit_code = -1
133+
134+
135+
@then("we see all command outputs")
136+
def step_see_all_command_outputs(context):
137+
"""Verify that all command outputs are present."""
138+
output = context.cmd_output.decode('utf-8')
139+
# Should contain output from all commands
140+
assert "first" in output or "from_c" in output, f"Expected 'first' or 'from_c' in output, but got: {output}"
141+
assert "second" in output or "from_command" in output, f"Expected 'second' or 'from_command' in output, but got: {output}"
142+
# For the 3-command test, also check for third
143+
if "third" in output or "result" in output:
144+
assert "third" in output, f"Expected 'third' in output for 3-command test, but got: {output}"

‎tests/test_main.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,31 @@ def test_execute_statements_does_not_split_inside_literals(executor):
563563
assert "a;b" in outputs[0]
564564

565565

566+
def test_command_and_file_both_run(tmpdir):
567+
"""psql runs both -c and -f when given together; -c must not exit first."""
568+
sql_file = tmpdir.join("script.sql")
569+
sql_file.write("select 2;")
570+
cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile")))
571+
cli.commands = ["select 1;"]
572+
cli.input_files = [str(sql_file)]
573+
with mock.patch.object(cli, "_execute_statements", return_value=True) as mock_exec:
574+
with pytest.raises(SystemExit) as e:
575+
cli.run_cli()
576+
assert e.value.code == 0
577+
assert [c[0][0] for c in mock_exec.call_args_list] == ["select 1;", "select 2;"]
578+
579+
580+
def test_command_mode_alone_still_exits(tmpdir):
581+
"""With no -f, the -c block still exits on its own."""
582+
cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile")))
583+
cli.commands = ["select 1;"]
584+
with mock.patch.object(cli, "_execute_statements", return_value=True) as mock_exec:
585+
with pytest.raises(SystemExit) as e:
586+
cli.run_cli()
587+
assert e.value.code == 0
588+
mock_exec.assert_called_once_with("select 1;")
589+
590+
566591
def test_file_mode_runs_statements(tmpdir):
567592
"""-f wiring: the file content goes through _execute_statements."""
568593
sql_file = tmpdir.join("script.sql")

0 commit comments

Comments
 (0)