From 3775c05df4d7af079a8c3d45c7402966fef084bc Mon Sep 17 00:00:00 2001 From: Amjith Ramanujam Date: Fri, 18 Apr 2025 22:15:52 -0700 Subject: [PATCH 1/7] Implement init-command similar to mycli. --- pgcli/main.py | 34 ++++++++++++++++++++++++++++++++++ pgcli/pgclirc | 8 ++++++++ 2 files changed, 42 insertions(+) diff --git a/pgcli/main.py b/pgcli/main.py index 0765efb7..4e8f4a76 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -1499,6 +1499,12 @@ def echo_via_pager(self, text, color=None): default=None, help="Write all queries & output into a file, in addition to the normal output destination.", ) +@click.option( + "--init-command", + "init_command", + type=str, + help="SQL statement to execute after connecting.", +) @click.argument("dbname", default=lambda: None, envvar="PGDATABASE", nargs=1) @click.argument("username", default=lambda: None, envvar="PGUSER", nargs=1) def cli( @@ -1525,6 +1531,7 @@ def cli( list_dsn, warn, ssh_tunnel: str, + init_command: str, log_file: str, ): if version: @@ -1682,6 +1689,33 @@ def echo_error(msg: str): # of conflicting sources echo_error(e.args[0]) + # Merge init-commands: global, DSN-specific, then CLI-provided + init_cmds = [] + # 1) Global init-commands + global_section = pgcli.config.get("init-commands", {}) + for _, val in global_section.items(): + if isinstance(val, (list, tuple)): + init_cmds.extend(val) + elif val: + init_cmds.append(val) + # 2) DSN-specific init-commands + if dsn: + alias_section = pgcli.config.get("alias_dsn.init-commands", {}) + if dsn in alias_section: + val = alias_section.get(dsn) + if isinstance(val, (list, tuple)): + init_cmds.extend(val) + elif val: + init_cmds.append(val) + # 3) CLI-provided init-command + if init_command: + init_cmds.append(init_command) + if init_cmds: + click.echo("Running init commands: %s" % "; ".join(init_cmds)) + for cmd in init_cmds: + # Execute each init command + list(pgcli.pgexecute.run(cmd)) + if list_databases: cur, headers, status = pgcli.pgexecute.full_databases() diff --git a/pgcli/pgclirc b/pgcli/pgclirc index be106107..77618e60 100644 --- a/pgcli/pgclirc +++ b/pgcli/pgclirc @@ -238,6 +238,14 @@ output.null = "#808080" [alias_dsn] # example_dsn = postgresql://[user[:password]@][netloc][:port][/dbname] +# Initial commands to execute when connecting to any database. +[init-commands] +# example = "SET search_path TO myschema" + +# Initial commands to execute when connecting to a DSN alias. +[alias_dsn.init-commands] +# example_dsn = "SET search_path TO otherschema; SET timezone TO 'UTC'" + # Format for number representation # for decimal "d" - 12345678, ",d" - 12,345,678 # for float "g" - 123456.78, ",g" - 123,456.78 From c021d6549b5a5b577c1b74a77d87ee284c459612 Mon Sep 17 00:00:00 2001 From: Amjith Ramanujam Date: Fri, 18 Apr 2025 22:29:46 -0700 Subject: [PATCH 2/7] Add tests for init-command. --- changelog.rst | 10 ++++ tests/test_init_commands_simple.py | 82 ++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+) create mode 100644 tests/test_init_commands_simple.py diff --git a/changelog.rst b/changelog.rst index dcf886a9..12345115 100644 --- a/changelog.rst +++ b/changelog.rst @@ -1,3 +1,13 @@ +Upcoming (TBD) +============== + +Features: +--------- +* Add support for `init-command` to run when the connection is established. + * Command line option `--init-command` + * Provide `init-command` in the config file + * Support dsn specific init-command in the config file + 4.3.0 (2025-03-22) ================== diff --git a/tests/test_init_commands_simple.py b/tests/test_init_commands_simple.py new file mode 100644 index 00000000..7bf0e37e --- /dev/null +++ b/tests/test_init_commands_simple.py @@ -0,0 +1,82 @@ +import pytest +from click.testing import CliRunner + +from pgcli.main import cli, PGCli + +@pytest.fixture +def dummy_exec(monkeypatch, tmp_path): + # Capture executed commands + # Isolate config directory for tests + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path)) + dummy_cmds = [] + class DummyExec: + def run(self, cmd): + # Ignore ping SELECT 1 commands used for exiting CLI + if cmd.strip().upper() == 'SELECT 1': + return [] + # Record init commands + dummy_cmds.append(cmd) + return [] + + def fake_connect(self, *args, **kwargs): + self.pgexecute = DummyExec() + + monkeypatch.setattr(PGCli, "connect", fake_connect) + return dummy_cmds + +def test_init_command_option(dummy_exec): + "Test that --init-command triggers execution of the command." + runner = CliRunner() + # Use a custom init command and --ping to exit the CLI after init commands + result = runner.invoke(cli, ["--init-command", "SELECT foo", "--ping", "db", "user"]) + assert result.exit_code == 0 + # Should print the init command + assert "Running init commands: SELECT foo" in result.output + # Should exit via ping + assert "PONG" in result.output + # DummyExec should have recorded only the init command + assert dummy_exec == ["SELECT foo"] + +def test_init_commands_from_config(dummy_exec, tmp_path): + """ + Test that init commands defined in the config file are executed on startup. + """ + # Create a temporary config file with init-commands + config_file = tmp_path / "pgclirc_test" + config_file.write_text( + "[init-commands]\n" + "first = SELECT foo;\n" + "second = SELECT bar;\n" + ) + + runner = CliRunner() + # Use --ping to exit the CLI after init commands + result = runner.invoke(cli, ["--pgclirc", str(config_file), "--ping", "testdb", "testuser"]) + assert result.exit_code == 0 + # Should print both init commands in order (note trailing semicolons cause double ';;') + assert "Running init commands: SELECT foo;; SELECT bar;" in result.output + # DummyExec should have recorded both commands + assert dummy_exec == ["SELECT foo;", "SELECT bar;"] + +def test_init_commands_option_and_config(dummy_exec, tmp_path): + """ + Test that CLI-provided init command is appended after config-defined commands. + """ + # Create a temporary config file with init-commands + config_file = tmp_path / "pgclirc_test" + config_file.write_text( + "[init-commands]\n" + "first = SELECT foo;\n" + ) + + runner = CliRunner() + # Use --ping to exit the CLI after init commands + result = runner.invoke( + cli, + ["--pgclirc", str(config_file), "--init-command", "SELECT baz;", "--ping", "testdb", "testuser"] + ) + assert result.exit_code == 0 + # Should print config command followed by CLI option (double ';' between commands) + assert "Running init commands: SELECT foo;; SELECT baz;" in result.output + # DummyExec should record both commands in order + assert dummy_exec == ["SELECT foo;", "SELECT baz;"] \ No newline at end of file From c97298cf8aaaa3e373983be3f2a30acc410fe6f3 Mon Sep 17 00:00:00 2001 From: Amjith Ramanujam Date: Fri, 18 Apr 2025 22:36:03 -0700 Subject: [PATCH 3/7] Update the default pgclirc. --- pgcli/pgclirc | 1 + 1 file changed, 1 insertion(+) diff --git a/pgcli/pgclirc b/pgcli/pgclirc index 77618e60..ab49362f 100644 --- a/pgcli/pgclirc +++ b/pgcli/pgclirc @@ -232,6 +232,7 @@ output.null = "#808080" # Named queries are queries you can execute by name. [named queries] +# ver = "SELECT version()" # Here's where you can provide a list of connection string aliases. # You can use it by passing the -D option. `pgcli -D example_dsn` From 1023f57ceeebedcbc6276066d0993445953f7c30 Mon Sep 17 00:00:00 2001 From: Amjith Ramanujam Date: Sat, 19 Apr 2025 17:29:18 -0700 Subject: [PATCH 4/7] Update tests. --- tests/test_init_commands_simple.py | 46 +++++++++++++++++++++--------- 1 file changed, 32 insertions(+), 14 deletions(-) diff --git a/tests/test_init_commands_simple.py b/tests/test_init_commands_simple.py index 7bf0e37e..70695dfe 100644 --- a/tests/test_init_commands_simple.py +++ b/tests/test_init_commands_simple.py @@ -3,32 +3,43 @@ from pgcli.main import cli, PGCli + @pytest.fixture def dummy_exec(monkeypatch, tmp_path): # Capture executed commands # Isolate config directory for tests monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path)) dummy_cmds = [] + class DummyExec: def run(self, cmd): # Ignore ping SELECT 1 commands used for exiting CLI - if cmd.strip().upper() == 'SELECT 1': + if cmd.strip().upper() == "SELECT 1": return [] # Record init commands dummy_cmds.append(cmd) return [] + def get_timezone(self): + return "UTC" + + def set_timezone(self, *args, **kwargs): + pass + def fake_connect(self, *args, **kwargs): self.pgexecute = DummyExec() monkeypatch.setattr(PGCli, "connect", fake_connect) return dummy_cmds + def test_init_command_option(dummy_exec): - "Test that --init-command triggers execution of the command." + "Test that --init-command triggers execution of the command." runner = CliRunner() # Use a custom init command and --ping to exit the CLI after init commands - result = runner.invoke(cli, ["--init-command", "SELECT foo", "--ping", "db", "user"]) + result = runner.invoke( + cli, ["--init-command", "SELECT foo", "--ping", "db", "user"] + ) assert result.exit_code == 0 # Should print the init command assert "Running init commands: SELECT foo" in result.output @@ -36,7 +47,8 @@ def test_init_command_option(dummy_exec): assert "PONG" in result.output # DummyExec should have recorded only the init command assert dummy_exec == ["SELECT foo"] - + + def test_init_commands_from_config(dummy_exec, tmp_path): """ Test that init commands defined in the config file are executed on startup. @@ -44,39 +56,45 @@ def test_init_commands_from_config(dummy_exec, tmp_path): # Create a temporary config file with init-commands config_file = tmp_path / "pgclirc_test" config_file.write_text( - "[init-commands]\n" - "first = SELECT foo;\n" - "second = SELECT bar;\n" + "[main]\n[init-commands]\nfirst = SELECT foo;\nsecond = SELECT bar;\n" ) runner = CliRunner() # Use --ping to exit the CLI after init commands - result = runner.invoke(cli, ["--pgclirc", str(config_file), "--ping", "testdb", "testuser"]) + result = runner.invoke( + cli, ["--pgclirc", str(config_file.absolute()), "--ping", "testdb", "user"] + ) assert result.exit_code == 0 # Should print both init commands in order (note trailing semicolons cause double ';;') assert "Running init commands: SELECT foo;; SELECT bar;" in result.output # DummyExec should have recorded both commands assert dummy_exec == ["SELECT foo;", "SELECT bar;"] + def test_init_commands_option_and_config(dummy_exec, tmp_path): """ Test that CLI-provided init command is appended after config-defined commands. """ # Create a temporary config file with init-commands config_file = tmp_path / "pgclirc_test" - config_file.write_text( - "[init-commands]\n" - "first = SELECT foo;\n" - ) + config_file.write_text("[main]\n [init-commands]\nfirst = SELECT foo;\n") runner = CliRunner() # Use --ping to exit the CLI after init commands result = runner.invoke( cli, - ["--pgclirc", str(config_file), "--init-command", "SELECT baz;", "--ping", "testdb", "testuser"] + [ + "--pgclirc", + str(config_file), + "--init-command", + "SELECT baz;", + "--ping", + "testdb", + "user", + ], ) assert result.exit_code == 0 # Should print config command followed by CLI option (double ';' between commands) assert "Running init commands: SELECT foo;; SELECT baz;" in result.output # DummyExec should record both commands in order - assert dummy_exec == ["SELECT foo;", "SELECT baz;"] \ No newline at end of file + assert dummy_exec == ["SELECT foo;", "SELECT baz;"] From d1ebd5ce952a04880fc397cbc42d16138a517b6a Mon Sep 17 00:00:00 2001 From: Amjith Ramanujam Date: Sun, 20 Apr 2025 09:14:02 -0700 Subject: [PATCH 5/7] Downgrade click version to 8.1.7 The newest click 8.1.8 is breaking the behave tests. --- pyproject.toml | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6c736022..04087114 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,8 +25,8 @@ urls = { Homepage = "https://pgcli.com" } requires-python = ">=3.9" dependencies = [ "pgspecial>=2.0.0", - "click >= 4.1", - "Pygments>=2.0", # Pygments has to be Capitalcased. WTF? + "click >= 4.1,<8.1.8", + "Pygments>=2.0", # Pygments has to be Capitalcased. WTF? # We still need to use pt-2 unless pt-3 released on Fedora32 # see: https://github.com/dbcli/pgcli/pull/1197 "prompt_toolkit>=2.0.6,<4.0.0", @@ -66,10 +66,7 @@ version = { attr = "pgcli.__version__" } find = { namespaces = false } [tool.setuptools.package-data] -pgcli = [ - "pgclirc", - "packages/pgliterals/pgliterals.json", -] +pgcli = ["pgclirc", "packages/pgliterals/pgliterals.json"] [tool.black] line-length = 88 From 5aa912c5b00467d65bd271b0ccb7a4a8d5746694 Mon Sep 17 00:00:00 2001 From: Amjith Ramanujam Date: Sun, 20 Apr 2025 09:14:33 -0700 Subject: [PATCH 6/7] Reenable the behave tests in CI. Click 8.1.8 was breaking them. --- .github/workflows/ci.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 227047db..6ea35faa 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -77,13 +77,13 @@ jobs: - name: Run unit tests run: coverage run --source pgcli -m pytest - # - name: Run integration tests - # env: - # PGUSER: postgres - # PGPASSWORD: postgres - # TERM: xterm + - name: Run integration tests + env: + PGUSER: postgres + PGPASSWORD: postgres + TERM: xterm - # run: behave tests/features --no-capture + run: behave tests/features --no-capture - name: Check changelog for ReST compliance run: docutils --halt=warning changelog.rst >/dev/null From a98163ecf82d920735a8289ec9269a53dbbf30fe Mon Sep 17 00:00:00 2001 From: Amjith Ramanujam Date: Sun, 20 Apr 2025 10:40:55 -0700 Subject: [PATCH 7/7] Fix the behave test expectation about PING/PONG test. --- tests/features/steps/basic_commands.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py index 00ac277f..2aaf2b29 100644 --- a/tests/features/steps/basic_commands.py +++ b/tests/features/steps/basic_commands.py @@ -36,7 +36,7 @@ def step_ping_database(context): def step_get_pong_response(context): # exit code 0 is implied by the presence of cmd_output here, which # is only set on a successful run. - assert context.cmd_output.strip() == b"PONG", f"Output was {context.cmd_output}" + assert b"PONG" in context.cmd_output.strip(), f"Output was {context.cmd_output}" @when("we run dbcli")