diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 227047dbc..6ea35faa1 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 diff --git a/changelog.rst b/changelog.rst index dcf886a9f..123451153 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/pgcli/main.py b/pgcli/main.py index 0765efb7f..4e8f4a768 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 be106107a..ab49362f9 100644 --- a/pgcli/pgclirc +++ b/pgcli/pgclirc @@ -232,12 +232,21 @@ 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` [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 diff --git a/pyproject.toml b/pyproject.toml index 6c7360224..04087114d 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 diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py index 00ac277f5..2aaf2b291 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") diff --git a/tests/test_init_commands_simple.py b/tests/test_init_commands_simple.py new file mode 100644 index 000000000..70695dfe8 --- /dev/null +++ b/tests/test_init_commands_simple.py @@ -0,0 +1,100 @@ +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 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." + 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( + "[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.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("[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", + "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;"]