refactor: Add command registration pattern to SqlQueryRunner (#767)

Summary:

This change introduces a command registration pattern to `SqlQueryRunner`, replacing the hardcoded if-else branches in `Console::readCommands` with a flexible, extensible command handling system.

## Why

The previous implementation had all command handling logic scattered in `Console.cpp` with a series of if-else branches checking command prefixes. This made it difficult to:
- Add new commands without modifying Console
- Extend SqlQueryRunner for different use cases

## What Changed

**New Command Infrastructure in SqlQueryRunner:**
- Added `CommandResult` struct with `handled`, `message`, `outcomes`, and `shouldExit` fields
- Added `CommandHandler` function type for consistent handler signatures
- Implemented prefix-based command matching with longest-match-first semantics

**Public API (SqlQueryRunner.h):**
- `registerDefaultCommands(options)` - Sets up built-in commands (exit, quit, help, savehistory, clearhistory, session)
- `handleCommand(command)` - Dispatches commands to registered handlers or SQL executor

Differential Revision: D90933335

Author: Shanyue Wan

axiom/cli/Console.cpp

@@ -54,6 +54,16 @@ void Console::initialize() {
 
   // Disable logging to stderr if not in debug mode.
   FLAGS_logtostderr = FLAGS_debug;
+
+  const SqlQueryRunner::RunOptions options{
+      .numWorkers = FLAGS_num_workers,
+      .numDrivers = FLAGS_num_drivers,
+      .splitTargetBytes = FLAGS_split_target_bytes,
+      .optimizerTraceFlags = FLAGS_optimizer_trace,
+      .debugMode = FLAGS_debug,
+      .historyPath = FLAGS_data_path + "/.history",
+  };
+  runner_.registerCommands(options);
 }
 
 void Console::run() {
@@ -334,20 +344,20 @@ void Console::readCommands(const std::string& prompt) {
       continue;
     }
 
-    if (command.starts_with("exit") || command.starts_with("quit")) {
+    // Try registered commands first, then execute SQL.
+    auto result = runner_.handleCommand(command);
+    for (const auto& outcome : result.outcomes) {
+      if (outcome.message.has_value()) {
+        std::cout << outcome.message.value() << std::endl;
+      }
+      if (!outcome.results.empty()) {
+        printResults(outcome.results);
+      }
+    }
+    if (result.shouldExit) {
       break;
     }
-
-    if (command.starts_with("help")) {
-      static const char* helpText =
-          "Axiom Interactive SQL\n\n"
-          "Type SQL and end with ';'.\n"
-          "To set a flag, type 'flag <gflag_name> = <value>;' Leave a space on either side of '='.\n\n"
-          "Useful flags:\n\n"
-          "num_workers - Make a distributed plan for this many workers. Runs it in-process with remote exchanges with serialization and passing data in memory. If num_workers is 1, makes single node plans without remote exchanges.\n\n"
-          "num_drivers - Specifies the parallelism for workers. This many threads per pipeline per worker.\n\n";
-
-      std::cout << helpText;
+    if (result.handled) {
       continue;
     }
 
@@ -398,25 +408,7 @@ void Console::readCommands(const std::string& prompt) {
       }
       continue;
     }
-
-    if (sscanf(command.c_str(), "session %ms = %ms", &flag, &value) == 2) {
-      std::cout << "Session '" << flag << "' set to '" << value << "'"
-                << std::endl;
-      runner_.sessionConfig()[std::string(flag)] = std::string(value);
-      continue;
-    }
-
-    if (command.starts_with("savehistory")) {
-      runner_.saveHistory(FLAGS_data_path + "/.history");
-      continue;
-    }
-
-    if (command.starts_with("clearhistory")) {
-      runner_.clearHistory();
-      continue;
-    }
-
-    runNoThrow(command);
   }
 }
+
 } // namespace axiom::sql

axiom/cli/SqlQueryRunner.cpp

@@ -85,6 +85,53 @@ std::vector<velox::RowVectorPtr> fetchResults(runner::LocalRunner& runner) {
   return results;
 }
 
+CommandResult exitHandler(const std::string&) {
+  return CommandResult{.handled = true, .outcomes = {}, .shouldExit = true};
+}
+
+CommandResult helpHandler(const std::string&) {
+  static const char* helpText =
+      "Axiom Interactive SQL\n\n"
+      "Type SQL and end with ';'.\n"
+      "To set a flag, type 'flag <gflag_name> = <value>;' Leave a space on either side of '='.\n\n"
+      "Useful flags:\n\n"
+      "num_workers - Make a distributed plan for this many workers. Runs it in-process with remote exchanges with serialization and passing data in memory. If num_workers is 1, makes single node plans without remote exchanges.\n\n"
+      "num_drivers - Specifies the parallelism for workers. This many threads per pipeline per worker.\n\n";
+  return CommandResult{
+      .handled = true, .outcomes = {{.message = helpText, .results = {}}}};
+}
+
+CommandResult sessionHandler(
+    const std::string& command,
+    std::unordered_map<std::string, std::string>& config) {
+  char* name = nullptr;
+  char* value = nullptr;
+  SCOPE_EXIT {
+    if (name != nullptr) {
+      free(name);
+    }
+    if (value != nullptr) {
+      free(value);
+    }
+  };
+
+  if (sscanf(command.c_str(), "session %ms = %ms", &name, &value) == 2) {
+    config[std::string(name)] = std::string(value);
+    return CommandResult{
+        .handled = true,
+        .outcomes = {
+            {.message = fmt::format(
+                 "Session '{}' set to '{}'",
+                 std::string(name),
+                 std::string(value)),
+             .results = {}}}};
+  }
+  return CommandResult{
+      .handled = true,
+      .outcomes = {
+          {.message = "Usage: session <name> = <value>", .results = {}}}};
+}
+
 } // namespace
 
 connector::TablePtr SqlQueryRunner::createTable(
@@ -376,4 +423,68 @@ std::vector<velox::RowVectorPtr> SqlQueryRunner::runSql(
   return results;
 }
 
+void SqlQueryRunner::registerCommands(const RunOptions& options) {
+  registerCommand("exit", exitHandler);
+  registerCommand("quit", exitHandler);
+  registerCommand("help", helpHandler);
+
+  registerCommand(
+      "savehistory",
+      [this, historyPath = options.historyPath](const std::string&) {
+        saveHistory(historyPath);
+        return CommandResult{
+            .handled = true,
+            .outcomes = {{.message = "History saved.", .results = {}}}};
+      });
+
+  registerCommand("clearhistory", [this](const std::string&) {
+    clearHistory();
+    return CommandResult{
+        .handled = true,
+        .outcomes = {{.message = "History cleared.", .results = {}}}};
+  });
+
+  registerCommand("session", [this](const std::string& command) {
+    return sessionHandler(command, config_);
+  });
+
+  setExecuteHandler([this, options](const std::string& command) {
+    try {
+      auto statements = parseMultiple(command, options);
+      std::vector<StatementOutcome> outcomes;
+      outcomes.reserve(statements.size());
+      for (const auto& statement : statements) {
+        auto sqlResult = run(*statement, options);
+        outcomes.push_back(
+            {.message = sqlResult.message,
+             .results = std::move(sqlResult.results)});
+      }
+      return CommandResult{.handled = true, .outcomes = std::move(outcomes)};
+    } catch (const std::exception& e) {
+      return CommandResult{
+          .handled = true, .outcomes = {{.message = e.what(), .results = {}}}};
+    }
+  });
+}
+
+CommandResult SqlQueryRunner::handleCommand(const std::string& command) const {
+  std::string bestMatch;
+  const CommandHandler* bestHandler = nullptr;
+
+  for (const auto& [prefix, handler] : commands_) {
+    if (command.starts_with(prefix)) {
+      if (prefix.size() > bestMatch.size()) {
+        bestMatch = prefix;
+        bestHandler = &handler;
+      }
+    }
+  }
+
+  if (bestHandler) {
+    return (*bestHandler)(command);
+  }
+
+  return executeHandler_(command);
+}
+
 } // namespace axiom::sql

axiom/cli/SqlQueryRunner.h

@@ -24,8 +24,36 @@
 
 namespace axiom::sql {
 
+/// Result of a single SQL statement execution.
+struct StatementOutcome {
+  /// Optional message to display to the user.
+  std::optional<std::string> message;
+
+  /// Optional SQL query results to display.
+  std::vector<facebook::velox::RowVectorPtr> results;
+};
+
+/// Result of executing a console command.
+struct CommandResult {
+  /// Whether the command was handled.
+  bool handled{false};
+
+  /// Results from executing one or more statements.
+  std::vector<StatementOutcome> outcomes;
+
+  /// Whether the command loop should exit.
+  bool shouldExit{false};
+};
+
+/// Function type for console command handlers.
+/// @param command The full command string entered by the user.
+/// @return CommandResult indicating how the command was processed.
+using CommandHandler = std::function<CommandResult(const std::string& command)>;
+
 class SqlQueryRunner {
  public:
+  virtual ~SqlQueryRunner() = default;
+
   /// @param initializeConnectors Lambda to call to initialize connectors and
   /// return a pair of default {connector ID, schema}. Takes a reference to the
   /// history to allow for loading from persistent storage.
@@ -50,6 +78,9 @@ class SqlQueryRunner {
 
     /// If true, EXPLAIN ANALYZE output includes custom operator stats.
     bool debugMode{false};
+
+    /// Path to save/load command history.
+    std::string historyPath;
   };
 
   /// Runs a single SQL statement and returns the result.
@@ -67,10 +98,33 @@ class SqlQueryRunner {
       std::string_view sql,
       const RunOptions& options);
 
-  std::unordered_map<std::string, std::string>& sessionConfig() {
-    return config_;
+  /// Registers the console commands (e.g. exit, quit, help,
+  /// savehistory, clearhistory).
+  virtual void registerCommands(const RunOptions& options);
+
+  /// Attempts to handle a command using registered handlers.
+  /// Commands are matched by prefix in order of longest match first.
+  /// If no handler matches, the execute handler is called.
+  /// @param command The full command string entered by the user.
+  /// @return CommandResult from the matched handler or the execute handler.
+  CommandResult handleCommand(const std::string& command) const;
+
+ protected:
+  /// Sets the execute handler called when no registered command matches.
+  /// This handler is required and executes SQL commands.
+  /// @param handler The handler to call for SQL execution.
+  void setExecuteHandler(CommandHandler handler) {
+    executeHandler_ = std::move(handler);
   }
 
+  /// Registers a command handler for a given command prefix.
+  /// @param prefix The command prefix to match (e.g., "help", "flag").
+  /// @param handler The handler function to call when the command matches.
+  void registerCommand(const std::string& prefix, CommandHandler handler) {
+    commands_[prefix] = std::move(handler);
+  }
+
+ private:
   void saveHistory(const std::string& path) {
     history_->saveToFile(path);
   }
@@ -79,7 +133,6 @@ class SqlQueryRunner {
     history_ = std::make_unique<facebook::axiom::optimizer::VeloxHistory>();
   }
 
- private:
   std::shared_ptr<facebook::velox::core::QueryCtx> newQuery(
       const RunOptions& options);
 
@@ -146,6 +199,8 @@ class SqlQueryRunner {
   std::unique_ptr<facebook::axiom::optimizer::VeloxHistory> history_;
   std::unique_ptr<presto::PrestoParser> prestoParser_;
   int32_t queryCounter_{0};
+  std::unordered_map<std::string, CommandHandler> commands_;
+  CommandHandler executeHandler_;
 };
 
 } // namespace axiom::sql