Skip to content

Commit 9d9a912

Browse files
authored
Merge pull request #1504 from dbcli/amjith/init-command
Implement init-command similar to mycli.
2 parents 4c6d169 + a98163e commit 9d9a912

7 files changed

Lines changed: 163 additions & 13 deletions

File tree

.github/workflows/ci.yml

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,13 @@ jobs:
7777
- name: Run unit tests
7878
run: coverage run --source pgcli -m pytest
7979

80-
# - name: Run integration tests
81-
# env:
82-
# PGUSER: postgres
83-
# PGPASSWORD: postgres
84-
# TERM: xterm
80+
- name: Run integration tests
81+
env:
82+
PGUSER: postgres
83+
PGPASSWORD: postgres
84+
TERM: xterm
8585

86-
# run: behave tests/features --no-capture
86+
run: behave tests/features --no-capture
8787

8888
- name: Check changelog for ReST compliance
8989
run: docutils --halt=warning changelog.rst >/dev/null

changelog.rst

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,13 @@
1+
Upcoming (TBD)
2+
==============
3+
4+
Features:
5+
---------
6+
* Add support for `init-command` to run when the connection is established.
7+
* Command line option `--init-command`
8+
* Provide `init-command` in the config file
9+
* Support dsn specific init-command in the config file
10+
111
4.3.0 (2025-03-22)
212
==================
313

pgcli/main.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1499,6 +1499,12 @@ def echo_via_pager(self, text, color=None):
14991499
default=None,
15001500
help="Write all queries & output into a file, in addition to the normal output destination.",
15011501
)
1502+
@click.option(
1503+
"--init-command",
1504+
"init_command",
1505+
type=str,
1506+
help="SQL statement to execute after connecting.",
1507+
)
15021508
@click.argument("dbname", default=lambda: None, envvar="PGDATABASE", nargs=1)
15031509
@click.argument("username", default=lambda: None, envvar="PGUSER", nargs=1)
15041510
def cli(
@@ -1525,6 +1531,7 @@ def cli(
15251531
list_dsn,
15261532
warn,
15271533
ssh_tunnel: str,
1534+
init_command: str,
15281535
log_file: str,
15291536
):
15301537
if version:
@@ -1682,6 +1689,33 @@ def echo_error(msg: str):
16821689
# of conflicting sources
16831690
echo_error(e.args[0])
16841691

1692+
# Merge init-commands: global, DSN-specific, then CLI-provided
1693+
init_cmds = []
1694+
# 1) Global init-commands
1695+
global_section = pgcli.config.get("init-commands", {})
1696+
for _, val in global_section.items():
1697+
if isinstance(val, (list, tuple)):
1698+
init_cmds.extend(val)
1699+
elif val:
1700+
init_cmds.append(val)
1701+
# 2) DSN-specific init-commands
1702+
if dsn:
1703+
alias_section = pgcli.config.get("alias_dsn.init-commands", {})
1704+
if dsn in alias_section:
1705+
val = alias_section.get(dsn)
1706+
if isinstance(val, (list, tuple)):
1707+
init_cmds.extend(val)
1708+
elif val:
1709+
init_cmds.append(val)
1710+
# 3) CLI-provided init-command
1711+
if init_command:
1712+
init_cmds.append(init_command)
1713+
if init_cmds:
1714+
click.echo("Running init commands: %s" % "; ".join(init_cmds))
1715+
for cmd in init_cmds:
1716+
# Execute each init command
1717+
list(pgcli.pgexecute.run(cmd))
1718+
16851719
if list_databases:
16861720
cur, headers, status = pgcli.pgexecute.full_databases()
16871721

pgcli/pgclirc

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,12 +232,21 @@ output.null = "#808080"
232232

233233
# Named queries are queries you can execute by name.
234234
[named queries]
235+
# ver = "SELECT version()"
235236

236237
# Here's where you can provide a list of connection string aliases.
237238
# You can use it by passing the -D option. `pgcli -D example_dsn`
238239
[alias_dsn]
239240
# example_dsn = postgresql://[user[:password]@][netloc][:port][/dbname]
240241

242+
# Initial commands to execute when connecting to any database.
243+
[init-commands]
244+
# example = "SET search_path TO myschema"
245+
246+
# Initial commands to execute when connecting to a DSN alias.
247+
[alias_dsn.init-commands]
248+
# example_dsn = "SET search_path TO otherschema; SET timezone TO 'UTC'"
249+
241250
# Format for number representation
242251
# for decimal "d" - 12345678, ",d" - 12,345,678
243252
# for float "g" - 123456.78, ",g" - 123,456.78

pyproject.toml

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@ urls = { Homepage = "https://pgcli.com" }
2525
requires-python = ">=3.9"
2626
dependencies = [
2727
"pgspecial>=2.0.0",
28-
"click >= 4.1",
29-
"Pygments>=2.0", # Pygments has to be Capitalcased. WTF?
28+
"click >= 4.1,<8.1.8",
29+
"Pygments>=2.0", # Pygments has to be Capitalcased. WTF?
3030
# We still need to use pt-2 unless pt-3 released on Fedora32
3131
# see: https://github.com/dbcli/pgcli/pull/1197
3232
"prompt_toolkit>=2.0.6,<4.0.0",
@@ -66,10 +66,7 @@ version = { attr = "pgcli.__version__" }
6666
find = { namespaces = false }
6767

6868
[tool.setuptools.package-data]
69-
pgcli = [
70-
"pgclirc",
71-
"packages/pgliterals/pgliterals.json",
72-
]
69+
pgcli = ["pgclirc", "packages/pgliterals/pgliterals.json"]
7370

7471
[tool.black]
7572
line-length = 88

tests/features/steps/basic_commands.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ def step_ping_database(context):
3636
def step_get_pong_response(context):
3737
# exit code 0 is implied by the presence of cmd_output here, which
3838
# is only set on a successful run.
39-
assert context.cmd_output.strip() == b"PONG", f"Output was {context.cmd_output}"
39+
assert b"PONG" in context.cmd_output.strip(), f"Output was {context.cmd_output}"
4040

4141

4242
@when("we run dbcli")

tests/test_init_commands_simple.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
import pytest
2+
from click.testing import CliRunner
3+
4+
from pgcli.main import cli, PGCli
5+
6+
7+
@pytest.fixture
8+
def dummy_exec(monkeypatch, tmp_path):
9+
# Capture executed commands
10+
# Isolate config directory for tests
11+
monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path))
12+
dummy_cmds = []
13+
14+
class DummyExec:
15+
def run(self, cmd):
16+
# Ignore ping SELECT 1 commands used for exiting CLI
17+
if cmd.strip().upper() == "SELECT 1":
18+
return []
19+
# Record init commands
20+
dummy_cmds.append(cmd)
21+
return []
22+
23+
def get_timezone(self):
24+
return "UTC"
25+
26+
def set_timezone(self, *args, **kwargs):
27+
pass
28+
29+
def fake_connect(self, *args, **kwargs):
30+
self.pgexecute = DummyExec()
31+
32+
monkeypatch.setattr(PGCli, "connect", fake_connect)
33+
return dummy_cmds
34+
35+
36+
def test_init_command_option(dummy_exec):
37+
"Test that --init-command triggers execution of the command."
38+
runner = CliRunner()
39+
# Use a custom init command and --ping to exit the CLI after init commands
40+
result = runner.invoke(
41+
cli, ["--init-command", "SELECT foo", "--ping", "db", "user"]
42+
)
43+
assert result.exit_code == 0
44+
# Should print the init command
45+
assert "Running init commands: SELECT foo" in result.output
46+
# Should exit via ping
47+
assert "PONG" in result.output
48+
# DummyExec should have recorded only the init command
49+
assert dummy_exec == ["SELECT foo"]
50+
51+
52+
def test_init_commands_from_config(dummy_exec, tmp_path):
53+
"""
54+
Test that init commands defined in the config file are executed on startup.
55+
"""
56+
# Create a temporary config file with init-commands
57+
config_file = tmp_path / "pgclirc_test"
58+
config_file.write_text(
59+
"[main]\n[init-commands]\nfirst = SELECT foo;\nsecond = SELECT bar;\n"
60+
)
61+
62+
runner = CliRunner()
63+
# Use --ping to exit the CLI after init commands
64+
result = runner.invoke(
65+
cli, ["--pgclirc", str(config_file.absolute()), "--ping", "testdb", "user"]
66+
)
67+
assert result.exit_code == 0
68+
# Should print both init commands in order (note trailing semicolons cause double ';;')
69+
assert "Running init commands: SELECT foo;; SELECT bar;" in result.output
70+
# DummyExec should have recorded both commands
71+
assert dummy_exec == ["SELECT foo;", "SELECT bar;"]
72+
73+
74+
def test_init_commands_option_and_config(dummy_exec, tmp_path):
75+
"""
76+
Test that CLI-provided init command is appended after config-defined commands.
77+
"""
78+
# Create a temporary config file with init-commands
79+
config_file = tmp_path / "pgclirc_test"
80+
config_file.write_text("[main]\n [init-commands]\nfirst = SELECT foo;\n")
81+
82+
runner = CliRunner()
83+
# Use --ping to exit the CLI after init commands
84+
result = runner.invoke(
85+
cli,
86+
[
87+
"--pgclirc",
88+
str(config_file),
89+
"--init-command",
90+
"SELECT baz;",
91+
"--ping",
92+
"testdb",
93+
"user",
94+
],
95+
)
96+
assert result.exit_code == 0
97+
# Should print config command followed by CLI option (double ';' between commands)
98+
assert "Running init commands: SELECT foo;; SELECT baz;" in result.output
99+
# DummyExec should record both commands in order
100+
assert dummy_exec == ["SELECT foo;", "SELECT baz;"]

0 commit comments

Comments
 (0)