Antalya 26.3: add --jwt-command to client

docs/en/interfaces/cli.md

@@ -836,6 +836,8 @@ All command-line options can be specified directly on the command line or as def
 | `-d [ --database ] <database>`   | Select the database to default to for this connection.                                                                                                                                                                                                                                                                             | The current database from the server settings (`default` by default)                                             |
 | `-h [ --host ] <host>`           | The hostname of the ClickHouse server to connect to. Can either be a hostname or an IPv4 or IPv6 address. Multiple hosts can be passed via multiple arguments.                                                                                                                                                                    | `localhost`                                                                                                      |
 | `--jwt <value>`                  | Use JSON Web Token (JWT) for authentication. <br/><br/>Server JWT authorization is only available in ClickHouse Cloud.                                                                                                                                                                                                            | -                                                                                                                |
+| `--jwt-command <command>`        | Shell command whose stdout is used as the JWT. Invoked at startup and on every (re)connect. See [`--jwt-command` details](#jwt-command-details) below.                                                                                                                                                                             | -                                                                                                                |
+| `--jwt-command-timeout <seconds>` | Timeout for `--jwt-command`. Also settable as `<jwt-command-timeout>` in the config file; CLI wins.                                                                                                                                                                                                                                  | `30`                                                                                                            |
 | `--login[=<mode>]`       | Authenticate via OAuth2. Bare `--login` (no `=<mode>`) triggers ClickHouse Cloud automatic login — the provider is inferred from the server. To authenticate against a custom OpenID Connect provider, supply a `mode` and `--oauth-credentials`: `--login=browser` runs the Authorization Code + PKCE flow (opens a browser), `--login=device` runs the Device Authorization flow (prints a URL and short code — no browser needed). | -                                                                                                                |
 | `--oauth-credentials <path>` | Path to an OAuth2 credentials JSON file (Google Cloud Console format). Required when using `--login=browser` or `--login=device` with a custom OpenID Connect provider. See [OAuth credentials file format](#oauth-credentials-file) below. Refresh tokens are cached in `~/.clickhouse-client/oauth_cache.json` (mode `0600`). | `~/.clickhouse-client/oauth_client.json`                                                                        |
 | `--no-warnings`                  | Disable showing warnings from `system.warnings` when the client connects to the server.                                                                                                                                                                                                                                            | -                                                                                                                |
@@ -880,6 +882,16 @@ The default path is `~/.clickhouse-client/oauth_client.json`. Override it with `
 
 After a successful login the obtained refresh token is cached in `~/.clickhouse-client/oauth_cache.json` (file mode `0600`). Subsequent runs reuse the cached token silently and only open the browser or print a device code when the refresh token has expired.
 
+### `--jwt-command` details {#jwt-command-details}
+
+The command is executed via `/bin/sh -c` and stdout is taken as the JWT (one trailing newline is stripped). Stderr is forwarded to the client's stderr; stdin is closed. The command runs at client startup and on every (re)connection to the server — the token is treated as opaque, so caching/refresh is the script's responsibility.
+
+```bash
+clickhouse-client --jwt-command "curl -sS https://idp.example/token | jq -r .access_token"
+```
+
+Cannot be combined with `--jwt`, `--login`, or a non-default `--user`. Non-zero exit, empty output, or exceeding `--jwt-command-timeout` (default `30`s, overridable via `<jwt-command-timeout>` in `~/.clickhouse-client/config.xml`) fails authentication. On timeout the entire helper subprocess tree is terminated.
+
 ### Query options {#command-line-options-query}
 
 | Option                          | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |

programs/client/Client.cpp

@@ -26,6 +26,7 @@
 #include <Interpreters/Context.h>
 
 #include <Client/JWTProvider.h>
+#include <Client/CommandJWTProvider.h>
 #include <Client/ClientBaseHelpers.h>
 #include <Client/OAuthLogin.h>
 
@@ -374,6 +375,17 @@ try
     }
 
 #if USE_JWT_CPP && USE_SSL
+    if (config().has("jwt-command") && !config().has("jwt"))
+    {
+        int timeout = config().getInt("jwt-command-timeout", DEFAULT_JWT_COMMAND_TIMEOUT_SECONDS);
+        if (timeout <= 0)
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "jwt-command-timeout must be positive, got {}", timeout);
+
+        auto provider = std::make_shared<CommandJWTProvider>(config().getString("jwt-command"), timeout);
+        config().setString("jwt", provider->getJWT());
+        jwt_provider = std::move(provider);
+    }
+
     if (config().getBool("cloud_oauth_pending", false) && !config().has("jwt"))
     {
         login();
@@ -742,6 +754,10 @@ void Client::printHelpMessage(const OptionsDescription & options_description)
 
 void Client::addExtraOptions(OptionsDescription & options_description)
 {
+    static const std::string jwt_command_timeout_help =
+        "Timeout in seconds for --jwt-command. Default: " + std::to_string(DEFAULT_JWT_COMMAND_TIMEOUT_SECONDS)
+        + ". Also configurable as <jwt-command-timeout> in the client config file.";
+
     /// Main commandline options related to client functionality and all parameters from Settings.
     options_description.main_description->add_options()
         ("config,c", po::value<std::string>(), "config-file path (another shorthand)")
@@ -756,6 +772,9 @@ void Client::addExtraOptions(OptionsDescription & options_description)
         ("ssh-key-passphrase", po::value<std::string>(), "Passphrase for the SSH private key specified by --ssh-key-file.")
         ("quota_key", po::value<std::string>(), "A string to differentiate quotas when the user have keyed quotas configured on server")
         ("jwt", po::value<std::string>(), "Use JWT for authentication")
+        ("jwt-command", po::value<std::string>(),
+            "Shell command whose stdout is used as the JWT. Invoked at startup and on every (re)connect.")
+        ("jwt-command-timeout", po::value<int>(), jwt_command_timeout_help.c_str())
         ("one-time-password", po::value<std::string>(), "Time-based one-time password (TOTP) for two-factor authentication")
         ("login", po::value<std::string>()->implicit_value(""),
             "Authenticate via OAuth2. Optional mode: 'browser' (auth-code + PKCE, opens browser) "
@@ -931,6 +950,27 @@ void Client::processOptions(
         config().setString("jwt", options["jwt"].as<std::string>());
         config().setString("user", "");
     }
+    if (options.contains("jwt-command-timeout"))
+        config().setInt("jwt-command-timeout", options["jwt-command-timeout"].as<int>());
+
+    if (options.contains("jwt-command"))
+    {
+#if USE_JWT_CPP && USE_SSL
+        if (options.contains("jwt"))
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "--jwt-command and --jwt cannot both be specified");
+        if (options.contains("login"))
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "--jwt-command and --login cannot both be specified");
+        if (!options["user"].defaulted())
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "User and JWT flags can't be specified together");
+
+        /// Defer execution to Client::main, after processConfig has loaded the XML config.
+        /// Reading config().getInt("jwt-command-timeout", ...) here would miss the XML value.
+        config().setString("jwt-command", options["jwt-command"].as<std::string>());
+        config().setString("user", "");
+#else
+        throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, "JWT is disabled, because ClickHouse is built without JWT or SSL support");
+#endif
+    }
     if (options.count("oauth-credentials") && !options.count("login"))
         throw Exception(
             ErrorCodes::BAD_ARGUMENTS,

src/Client/CommandJWTProvider.cpp

@@ -0,0 +1,128 @@
+#include <config.h>
+
+#if USE_JWT_CPP && USE_SSL
+#include <Client/CommandJWTProvider.h>
+
+#include <Common/Exception.h>
+#include <Common/ShellCommand.h>
+#include <IO/ReadHelpers.h>
+#include <IO/WriteBufferFromOStream.h>
+#include <IO/copyData.h>
+#include <base/scope_guard.h>
+
+#include <csignal>
+#include <chrono>
+#include <condition_variable>
+#include <iostream>
+#include <mutex>
+#include <thread>
+#include <utility>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int AUTHENTICATION_FAILED;
+}
+
+CommandJWTProvider::CommandJWTProvider(std::string command_, int timeout_seconds_)
+    : JWTProvider(/*auth_url=*/"", /*client_id=*/"", /*audience=*/"", std::cout, std::cerr)
+    , command(std::move(command_))
+    , timeout_seconds(timeout_seconds_)
+{
+}
+
+std::string CommandJWTProvider::getJWT()
+{
+    ShellCommand::Config config(command);
+    config.new_process_group = true;  // so the watchdog can kill the whole tree, not just /bin/sh
+    auto child = ShellCommand::execute(config);
+    child->in.close();  // we don't write to the script's stdin; close so reads see EOF
+    const pid_t pid = child->getPid();
+
+    std::mutex mutex;
+    std::condition_variable cv;
+    bool finished = false;
+    std::atomic<bool> timed_out{false};
+
+    /// Default-construct the threads first and install the cleanup guard before assigning,
+    /// so a thread-constructor failure mid-assignment cannot leave a joinable thread that
+    /// would call std::terminate on destruction.
+    std::thread watchdog;
+    std::thread stderr_forwarder;
+    SCOPE_EXIT({
+        {
+            std::lock_guard lock(mutex);
+            finished = true;
+        }
+        cv.notify_all();
+        if (stderr_forwarder.joinable()) stderr_forwarder.join();
+        if (watchdog.joinable()) watchdog.join();
+    });
+
+    watchdog = std::thread([&, pid]()
+    {
+        std::unique_lock lock(mutex);
+        if (!cv.wait_for(lock, std::chrono::seconds(timeout_seconds), [&]{ return finished; }))
+        {
+            timed_out = true;
+            ::kill(-pid, SIGKILL);
+        }
+    });
+
+    /// Drain stderr on a separate thread so the child doesn't block on a full pipe.
+    stderr_forwarder = std::thread([&child]()
+    {
+        try
+        {
+            WriteBufferFromOStream wb(std::cerr);
+            copyData(child->err, wb);
+            wb.finalize();
+        }
+        catch (...) {}
+    });
+
+    std::string token;
+    readStringUntilEOF(token, child->out);
+
+    /// Drain stderr fully before tryWait, since tryWait closes child->err and reading
+    /// a buffer whose fd has just been closed from another thread is UB.
+    stderr_forwarder.join();
+
+    /// Cancel the watchdog before tryWait. After tryWait reaps the child, the kernel
+    /// may recycle the pid; if the watchdog then fires kill(-pid, ...) it could hit an
+    /// unrelated process group.
+    {
+        std::lock_guard lock(mutex);
+        finished = true;
+    }
+    cv.notify_all();
+    watchdog.join();
+
+    /// Reap with a catch: on timeout the child is signaled, and we want our own
+    /// error message rather than the noisy CHILD_WAS_NOT_EXITED_NORMALLY one.
+    int retcode = 0;
+    try { retcode = child->tryWait(); }
+    catch (...) { if (!timed_out.load()) throw; }
+
+    if (timed_out.load())
+        throw Exception(ErrorCodes::AUTHENTICATION_FAILED,
+            "--jwt-command timed out after {} seconds", timeout_seconds);
+
+    if (retcode != 0)
+        throw Exception(ErrorCodes::AUTHENTICATION_FAILED,
+            "--jwt-command exited with non-zero status {}", retcode);
+
+    if (!token.empty() && token.back() == '\n')
+        token.pop_back();
+
+    if (token.empty())
+        throw Exception(ErrorCodes::AUTHENTICATION_FAILED, "--jwt-command produced empty output");
+
+    return token;
+}
+
+}
+
+#endif

src/Client/CommandJWTProvider.h

@@ -0,0 +1,33 @@
+#pragma once
+
+#include <string>
+
+namespace DB
+{
+
+inline constexpr int DEFAULT_JWT_COMMAND_TIMEOUT_SECONDS = 30;
+
+}
+
+#if USE_JWT_CPP && USE_SSL
+
+#include <Client/JWTProvider.h>
+
+namespace DB
+{
+
+class CommandJWTProvider : public JWTProvider
+{
+public:
+    CommandJWTProvider(std::string command_, int timeout_seconds_);
+
+    std::string getJWT() override;
+
+private:
+    std::string command;
+    int timeout_seconds;
+};
+
+}
+
+#endif

src/Client/Connection.cpp

@@ -146,6 +146,11 @@ void Connection::connect(const ConnectionTimeouts & timeouts)
     /// if connection was broken it is necessary to cancel it before reconnecting
     disconnect();
 
+#if USE_JWT_CPP && USE_SSL
+    if (jwt_provider)
+        jwt = jwt_provider->getJWT();
+#endif
+
     ProfileEvents::increment(ProfileEvents::DistributedConnectionConnectCount);
     try
     {
@@ -848,23 +853,6 @@ void Connection::sendQuery(
         client_info = &new_client_info;
     }
 
-#if USE_JWT_CPP && USE_SSL
-    if (jwt_provider && !jwt.empty())
-    {
-        if (JWTProvider::getJwtExpiry(jwt) < (Poco::Timestamp() + Poco::Timespan(30, 0)))
-        {
-            String new_jwt = jwt_provider->getJWT();
-            if (!new_jwt.empty())
-            {
-                jwt = new_jwt;
-                // We have a new token, so we need to reconnect.
-                // The current connection is still using the old token.
-                disconnect();
-            }
-        }
-    }
-#endif
-
     if (!connected)
         connect(timeouts);
 

src/Common/ShellCommand.cpp

@@ -27,6 +27,7 @@ namespace
         CANNOT_EXEC                 = 0x55555558,
         CANNOT_DUP_READ_DESCRIPTOR  = 0x55555559,
         CANNOT_DUP_WRITE_DESCRIPTOR = 0x55555560,
+        CANNOT_SETPGID              = 0x55555561,
     };
 }
 
@@ -218,6 +219,9 @@ std::unique_ptr<ShellCommand> ShellCommand::executeImpl(
         sigprocmask(0, nullptr, &mask); // NOLINT(concurrency-mt-unsafe)
         sigprocmask(SIG_UNBLOCK, &mask, nullptr); // NOLINT(concurrency-mt-unsafe)
 
+        if (config.new_process_group && setpgid(0, 0) != 0)
+            _exit(static_cast<int>(ReturnCodes::CANNOT_SETPGID));
+
         execv(filename, argv);
         /// If the process is running, then `execv` does not return here.
 
@@ -385,6 +389,8 @@ void ShellCommand::handleProcessRetcode(int retcode) const
                 throw Exception(ErrorCodes::CANNOT_CREATE_CHILD_PROCESS, "Cannot dup2 read descriptor of child process");
             case static_cast<int>(ReturnCodes::CANNOT_DUP_WRITE_DESCRIPTOR):
                 throw Exception(ErrorCodes::CANNOT_CREATE_CHILD_PROCESS, "Cannot dup2 write descriptor of child process");
+            case static_cast<int>(ReturnCodes::CANNOT_SETPGID):
+                throw Exception(ErrorCodes::CANNOT_CREATE_CHILD_PROCESS, "Cannot setpgid in child process");
             default:
                 throw Exception(ErrorCodes::CHILD_WAS_NOT_EXITED_NORMALLY, "Child process was exited with return code {}", toString(retcode));
         }

src/Common/ShellCommand.h

@@ -65,6 +65,10 @@ class ShellCommand final
 
         bool pipe_stdin_only = false;
 
+        /// Put the child in its own process group, so that a single `kill(-pid, ...)`
+        /// from the parent terminates the entire subprocess tree.
+        bool new_process_group = false;
+
         DestructorStrategy terminate_in_destructor_strategy = DestructorStrategy(false, 0);
     };
 

tests/queries/0_stateless/04206_jwt_command.reference

@@ -0,0 +1,25 @@
+Test 1: --jwt and --jwt-command together should give BAD_ARGUMENTS
+OK
+Test 2: --jwt-command with non-default --user should give BAD_ARGUMENTS
+OK
+Test 3: --jwt-command with --login should give BAD_ARGUMENTS
+OK
+Test 4: --jwt-command with empty stdout should fail with AUTHENTICATION_FAILED
+OK
+Test 5: --jwt-command exiting with non-zero status should fail with AUTHENTICATION_FAILED
+OK
+Test 6: --jwt-command stderr should be forwarded to client stderr
+OK
+Test 7: --jwt-command-timeout kills a hanging script
+OK
+Test 8: --jwt-command-timeout=0 should be rejected
+OK
+Test 9: --jwt-command is actually executed
+OK
+Test 10: --jwt-command-timeout from XML config file takes effect
+OK
+Test 11: stdin-reading script completes promptly (stdin is closed)
+OK
+Test 12: CLI --jwt-command-timeout overrides XML config
+OK
+All tests completed