Skip to content

Commit 61cff59

Browse files
Shanyue Wanfacebook-github-bot
Shanyue Wan
authored and
facebook-github-bot
committed
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
1 parent 56691e9 commit 61cff59

File tree

4 files changed

+197
-45
lines changed

4 files changed

+197
-45
lines changed

‎axiom/cli/Console.cpp‎

Lines changed: 25 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,21 @@ void Console::initialize() {
5454

5555
// Disable logging to stderr if not in debug mode.
5656
FLAGS_logtostderr = FLAGS_debug;
57+
58+
const SqlQueryRunner::RunOptions options{
59+
.numWorkers = FLAGS_num_workers,
60+
.numDrivers = FLAGS_num_drivers,
61+
.splitTargetBytes = FLAGS_split_target_bytes,
62+
.optimizerTraceFlags = FLAGS_optimizer_trace,
63+
.debugMode = FLAGS_debug,
64+
.historyPath = FLAGS_data_path + "/.history",
65+
};
66+
runner_.registerCommands(options);
5767
}
5868

5969
void Console::run() {
6070
if (!FLAGS_query.empty()) {
61-
runNoThrow(FLAGS_query, false);
71+
runNoThrow(FLAGS_query);
6272
} else {
6373
std::cout << "Axiom SQL. Type statement and end with ;.\n"
6474
"flag name = value; sets a gflag.\n"
@@ -215,7 +225,7 @@ int32_t printResults(const std::vector<RowVectorPtr>& results) {
215225
}
216226
} // namespace
217227

218-
void Console::runNoThrow(std::string_view sql, bool isInteractive) {
228+
void Console::runNoThrow(std::string_view sql) {
219229
const SqlQueryRunner::RunOptions options{
220230
.numWorkers = FLAGS_num_workers,
221231
.numDrivers = FLAGS_num_drivers,
@@ -245,11 +255,6 @@ void Console::runNoThrow(std::string_view sql, bool isInteractive) {
245255
} else {
246256
printResults(result.results);
247257
}
248-
249-
if (isInteractive) {
250-
// In interactive mode, show per-statement timing.
251-
std::cout << statementTiming.toString() << std::endl;
252-
}
253258
} catch (std::exception& e) {
254259
std::cerr << "Query failed: " << e.what() << std::endl;
255260
// Stop executing remaining statements in this block.
@@ -334,20 +339,20 @@ void Console::readCommands(const std::string& prompt) {
334339
continue;
335340
}
336341

337-
if (command.starts_with("exit") || command.starts_with("quit")) {
342+
// Try registered commands first, then execute SQL.
343+
auto result = runner_.handleCommand(command);
344+
for (const auto& outcome : result.outcomes) {
345+
if (outcome.message.has_value()) {
346+
std::cout << outcome.message.value() << std::endl;
347+
}
348+
if (!outcome.results.empty()) {
349+
printResults(outcome.results);
350+
}
351+
}
352+
if (result.shouldExit) {
338353
break;
339354
}
340-
341-
if (command.starts_with("help")) {
342-
static const char* helpText =
343-
"Axiom Interactive SQL\n\n"
344-
"Type SQL and end with ';'.\n"
345-
"To set a flag, type 'flag <gflag_name> = <value>;' Leave a space on either side of '='.\n\n"
346-
"Useful flags:\n\n"
347-
"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"
348-
"num_drivers - Specifies the parallelism for workers. This many threads per pipeline per worker.\n\n";
349-
350-
std::cout << helpText;
355+
if (result.handled) {
351356
continue;
352357
}
353358

@@ -398,25 +403,7 @@ void Console::readCommands(const std::string& prompt) {
398403
}
399404
continue;
400405
}
401-
402-
if (sscanf(command.c_str(), "session %ms = %ms", &flag, &value) == 2) {
403-
std::cout << "Session '" << flag << "' set to '" << value << "'"
404-
<< std::endl;
405-
runner_.sessionConfig()[std::string(flag)] = std::string(value);
406-
continue;
407-
}
408-
409-
if (command.starts_with("savehistory")) {
410-
runner_.saveHistory(FLAGS_data_path + "/.history");
411-
continue;
412-
}
413-
414-
if (command.starts_with("clearhistory")) {
415-
runner_.clearHistory();
416-
continue;
417-
}
418-
419-
runNoThrow(command);
420406
}
421407
}
408+
422409
} // namespace axiom::sql

‎axiom/cli/Console.h‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,11 @@ class Console {
3434
void run();
3535

3636
private:
37-
// Executes SQL and prints results, catching any exceptions.
37+
// Executes SQL and prints results, catching any exceptions. Shows timing
38+
// after each statement for multi-statement queries.
3839
// @param sql The SQL text to execute, which may contain multiple
3940
// semicolon-separated statements.
40-
// @param isInteractive If true, shows timing after each statement for
41-
// multi-statement queries.
42-
void runNoThrow(std::string_view sql, bool isInteractive = true);
41+
void runNoThrow(std::string_view sql);
4342

4443
// Reads and executes commands from standard input in interactive mode.
4544
void readCommands(const std::string& prompt);

‎axiom/cli/SqlQueryRunner.cpp‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,53 @@ std::vector<velox::RowVectorPtr> fetchResults(runner::LocalRunner& runner) {
8585
return results;
8686
}
8787

88+
CommandResult exitHandler(const std::string&) {
89+
return CommandResult{.handled = true, .outcomes = {}, .shouldExit = true};
90+
}
91+
92+
CommandResult helpHandler(const std::string&) {
93+
static const char* helpText =
94+
"Axiom Interactive SQL\n\n"
95+
"Type SQL and end with ';'.\n"
96+
"To set a flag, type 'flag <gflag_name> = <value>;' Leave a space on either side of '='.\n\n"
97+
"Useful flags:\n\n"
98+
"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"
99+
"num_drivers - Specifies the parallelism for workers. This many threads per pipeline per worker.\n\n";
100+
return CommandResult{
101+
.handled = true, .outcomes = {{.message = helpText, .results = {}}}};
102+
}
103+
104+
CommandResult sessionHandler(
105+
const std::string& command,
106+
std::unordered_map<std::string, std::string>& config) {
107+
char* name = nullptr;
108+
char* value = nullptr;
109+
SCOPE_EXIT {
110+
if (name != nullptr) {
111+
free(name);
112+
}
113+
if (value != nullptr) {
114+
free(value);
115+
}
116+
};
117+
118+
if (sscanf(command.c_str(), "session %ms = %ms", &name, &value) == 2) {
119+
config[std::string(name)] = std::string(value);
120+
return CommandResult{
121+
.handled = true,
122+
.outcomes = {
123+
{.message = fmt::format(
124+
"Session '{}' set to '{}'",
125+
std::string(name),
126+
std::string(value)),
127+
.results = {}}}};
128+
}
129+
return CommandResult{
130+
.handled = true,
131+
.outcomes = {
132+
{.message = "Usage: session <name> = <value>", .results = {}}}};
133+
}
134+
88135
} // namespace
89136

90137
connector::TablePtr SqlQueryRunner::createTable(
@@ -139,6 +186,70 @@ std::vector<presto::SqlStatementPtr> SqlQueryRunner::parseMultiple(
139186
return prestoParser_->parseMultiple(sql, /*enableTracing=*/options.debugMode);
140187
}
141188

189+
void SqlQueryRunner::registerCommands(const RunOptions& options) {
190+
registerCommand("exit", exitHandler);
191+
registerCommand("quit", exitHandler);
192+
registerCommand("help", helpHandler);
193+
194+
registerCommand(
195+
"savehistory",
196+
[this, historyPath = options.historyPath](const std::string&) {
197+
saveHistory(historyPath);
198+
return CommandResult{
199+
.handled = true,
200+
.outcomes = {{.message = "History saved.", .results = {}}}};
201+
});
202+
203+
registerCommand("clearhistory", [this](const std::string&) {
204+
clearHistory();
205+
return CommandResult{
206+
.handled = true,
207+
.outcomes = {{.message = "History cleared.", .results = {}}}};
208+
});
209+
210+
registerCommand("session", [this](const std::string& command) {
211+
return sessionHandler(command, config_);
212+
});
213+
214+
setExecuteHandler([this, options](const std::string& command) {
215+
std::vector<StatementOutcome> outcomes;
216+
try {
217+
auto statements = parseMultiple(command, options);
218+
outcomes.reserve(statements.size());
219+
for (const auto& statement : statements) {
220+
auto sqlResult = run(*statement, options);
221+
outcomes.push_back(
222+
{.message = sqlResult.message,
223+
.results = std::move(sqlResult.results)});
224+
}
225+
return CommandResult{.handled = true, .outcomes = std::move(outcomes)};
226+
} catch (const std::exception& e) {
227+
outcomes.push_back({.message = e.what(), .results = {}});
228+
return CommandResult{.handled = true, .outcomes = std::move(outcomes)};
229+
}
230+
});
231+
}
232+
233+
CommandResult SqlQueryRunner::handleCommand(const std::string& command) const {
234+
std::string bestMatch;
235+
const CommandHandler* bestHandler = nullptr;
236+
237+
for (const auto& [prefix, handler] : commands_) {
238+
if (command.starts_with(prefix)) {
239+
if (prefix.size() > bestMatch.size()) {
240+
bestMatch = prefix;
241+
bestHandler = &handler;
242+
}
243+
}
244+
}
245+
246+
if (bestHandler) {
247+
return (*bestHandler)(command);
248+
}
249+
250+
return executeHandler_(command);
251+
}
252+
142253
SqlQueryRunner::SqlResult SqlQueryRunner::run(
143254
const presto::SqlStatement& sqlStatement,
144255
const RunOptions& options) {

‎axiom/cli/SqlQueryRunner.h‎

Lines changed: 58 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,36 @@
2424

2525
namespace axiom::sql {
2626

27+
/// Result of a single SQL statement execution.
28+
struct StatementOutcome {
29+
/// Optional message to display to the user.
30+
std::optional<std::string> message;
31+
32+
/// Optional SQL query results to display.
33+
std::vector<facebook::velox::RowVectorPtr> results;
34+
};
35+
36+
/// Result of executing a console command.
37+
struct CommandResult {
38+
/// Whether the command was handled.
39+
bool handled{false};
40+
41+
/// Results from executing one or more statements.
42+
std::vector<StatementOutcome> outcomes;
43+
44+
/// Whether the command loop should exit.
45+
bool shouldExit{false};
46+
};
47+
48+
/// Function type for console command handlers.
49+
/// @param command The full command string entered by the user.
50+
/// @return CommandResult indicating how the command was processed.
51+
using CommandHandler = std::function<CommandResult(const std::string& command)>;
52+
2753
class SqlQueryRunner {
2854
public:
55+
virtual ~SqlQueryRunner() = default;
56+
2957
/// @param initializeConnectors Lambda to call to initialize connectors and
3058
/// return a pair of default {connector ID, schema}. Takes a reference to the
3159
/// history to allow for loading from persistent storage.
@@ -50,6 +78,9 @@ class SqlQueryRunner {
5078

5179
/// If true, EXPLAIN ANALYZE output includes custom operator stats.
5280
bool debugMode{false};
81+
82+
/// Path to save/load command history.
83+
std::string historyPath;
5384
};
5485

5586
/// Runs a single SQL statement and returns the result.
@@ -67,10 +98,33 @@ class SqlQueryRunner {
6798
std::string_view sql,
6899
const RunOptions& options);
69100

70-
std::unordered_map<std::string, std::string>& sessionConfig() {
71-
return config_;
101+
/// Registers the console commands (e.g. exit, quit, help,
102+
/// savehistory, clearhistory).
103+
virtual void registerCommands(const RunOptions& options);
104+
105+
/// Attempts to handle a command using registered handlers.
106+
/// Commands are matched by prefix in order of longest match first.
107+
/// If no handler matches, the execute handler is called.
108+
/// @param command The full command string entered by the user.
109+
/// @return CommandResult from the matched handler or the execute handler.
110+
CommandResult handleCommand(const std::string& command) const;
111+
112+
protected:
113+
/// Sets the execute handler called when no registered command matches.
114+
/// This handler is required and executes SQL commands.
115+
/// @param handler The handler to call for SQL execution.
116+
void setExecuteHandler(CommandHandler handler) {
117+
executeHandler_ = std::move(handler);
72118
}
73119

120+
/// Registers a command handler for a given command prefix.
121+
/// @param prefix The command prefix to match (e.g., "help", "flag").
122+
/// @param handler The handler function to call when the command matches.
123+
void registerCommand(const std::string& prefix, CommandHandler handler) {
124+
commands_[prefix] = std::move(handler);
125+
}
126+
127+
private:
74128
void saveHistory(const std::string& path) {
75129
history_->saveToFile(path);
76130
}
@@ -79,7 +133,6 @@ class SqlQueryRunner {
79133
history_ = std::make_unique<facebook::axiom::optimizer::VeloxHistory>();
80134
}
81135

82-
private:
83136
std::shared_ptr<facebook::velox::core::QueryCtx> newQuery(
84137
const RunOptions& options);
85138

@@ -146,6 +199,8 @@ class SqlQueryRunner {
146199
std::unique_ptr<facebook::axiom::optimizer::VeloxHistory> history_;
147200
std::unique_ptr<presto::PrestoParser> prestoParser_;
148201
int32_t queryCounter_{0};
202+
std::unordered_map<std::string, CommandHandler> commands_;
203+
CommandHandler executeHandler_;
149204
};
150205

151206
} // namespace axiom::sql

0 commit comments

Comments
 (0)