Skip to content

Commit 3775c05

Browse files
committed
Implement init-command similar to mycli.
1 parent 4c6d169 commit 3775c05

File tree

2 files changed

+42
-0
lines changed

2 files changed

+42
-0
lines changed

pgcli/main.py

+34
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

+8
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,14 @@ output.null = "#808080"
238238
[alias_dsn]
239239
# example_dsn = postgresql://[user[:password]@][netloc][:port][/dbname]
240240

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

0 commit comments

Comments
 (0)