refactor: Allow SqlQueryRunner to be extended (#767)

Summary:

Create a virtual class to let SqlRunner to be extended.

Differential Revision: D90933335

Author: Shanyue Wan

axiom/cli/AxiomSql.cpp

@@ -128,7 +128,7 @@ int main(int argc, char** argv) {
 
   facebook::axiom::Connectors connectors{FLAGS_data_path, FLAGS_data_format};
 
-  axiom::sql::SqlQueryRunner runner;
+  axiom::sql::AxiomSqlQueryRunner runner;
   runner.initialize([&](auto& history) {
     return std::make_pair(connectors.initialize(history), std::nullopt);
   });

axiom/cli/SqlQueryRunner.cpp

@@ -35,7 +35,7 @@ using namespace facebook::axiom;
 
 namespace axiom::sql {
 
-void SqlQueryRunner::initialize(
+void AxiomSqlQueryRunner::initialize(
     const std::function<std::pair<std::string, std::optional<std::string>>(
         optimizer::VeloxHistory& history)>& initializeConnectors) {
   static folly::once_flag kInitialized;
@@ -87,7 +87,7 @@ std::vector<velox::RowVectorPtr> fetchResults(runner::LocalRunner& runner) {
 
 } // namespace
 
-connector::TablePtr SqlQueryRunner::createTable(
+connector::TablePtr AxiomSqlQueryRunner::createTable(
     const presto::CreateTableAsSelectStatement& statement) {
   auto metadata =
       connector::ConnectorMetadata::metadata(statement.connectorId());
@@ -103,7 +103,7 @@ connector::TablePtr SqlQueryRunner::createTable(
       session, statement.tableName(), statement.tableSchema(), options);
 }
 
-std::string SqlQueryRunner::dropTable(
+std::string AxiomSqlQueryRunner::dropTable(
     const presto::DropTableStatement& statement) {
   auto metadata =
       connector::ConnectorMetadata::metadata(statement.connectorId());
@@ -121,7 +121,7 @@ std::string SqlQueryRunner::dropTable(
   }
 }
 
-SqlQueryRunner::SqlResult SqlQueryRunner::run(
+AxiomSqlQueryRunner::SqlResult AxiomSqlQueryRunner::run(
     std::string_view sql,
     const RunOptions& options) {
   auto statements = parseMultiple(sql, options);
@@ -133,13 +133,13 @@ SqlQueryRunner::SqlResult SqlQueryRunner::run(
   return run(*statements[0], options);
 }
 
-std::vector<presto::SqlStatementPtr> SqlQueryRunner::parseMultiple(
+std::vector<presto::SqlStatementPtr> AxiomSqlQueryRunner::parseMultiple(
     std::string_view sql,
     const RunOptions& options) {
   return prestoParser_->parseMultiple(sql, /*enableTracing=*/options.debugMode);
 }
 
-SqlQueryRunner::SqlResult SqlQueryRunner::run(
+AxiomSqlQueryRunner::SqlResult AxiomSqlQueryRunner::run(
     const presto::SqlStatement& sqlStatement,
     const RunOptions& options) {
   if (sqlStatement.isExplain()) {
@@ -198,7 +198,7 @@ SqlQueryRunner::SqlResult SqlQueryRunner::run(
   return {.results = runSql(logicalPlan, options)};
 }
 
-std::shared_ptr<velox::core::QueryCtx> SqlQueryRunner::newQuery(
+std::shared_ptr<velox::core::QueryCtx> AxiomSqlQueryRunner::newQuery(
     const RunOptions& options) {
   ++queryCounter_;
 
@@ -216,7 +216,7 @@ std::shared_ptr<velox::core::QueryCtx> SqlQueryRunner::newQuery(
       fmt::format("query_{}", queryCounter_));
 }
 
-std::string SqlQueryRunner::runExplain(
+std::string AxiomSqlQueryRunner::runExplain(
     const logical_plan::LogicalPlanNodePtr& logicalPlan,
     presto::ExplainStatement::Type type,
     const RunOptions& options) {
@@ -273,7 +273,7 @@ std::string printPlanWithStats(
 }
 } // namespace
 
-std::string SqlQueryRunner::runExplainAnalyze(
+std::string AxiomSqlQueryRunner::runExplainAnalyze(
     const logical_plan::LogicalPlanNodePtr& logicalPlan,
     const RunOptions& options) {
   auto queryCtx = newQuery(options);
@@ -293,7 +293,7 @@ std::string SqlQueryRunner::runExplainAnalyze(
   return out.str();
 }
 
-optimizer::PlanAndStats SqlQueryRunner::optimize(
+optimizer::PlanAndStats AxiomSqlQueryRunner::optimize(
     const logical_plan::LogicalPlanNodePtr& logicalPlan,
     const std::shared_ptr<velox::core::QueryCtx>& queryCtx,
     const RunOptions& options,
@@ -339,7 +339,7 @@ optimizer::PlanAndStats SqlQueryRunner::optimize(
   return optimization.toVeloxPlan(best->op);
 }
 
-std::shared_ptr<runner::LocalRunner> SqlQueryRunner::makeLocalRunner(
+std::shared_ptr<runner::LocalRunner> AxiomSqlQueryRunner::makeLocalRunner(
     optimizer::PlanAndStats& planAndStats,
     const std::shared_ptr<velox::core::QueryCtx>& queryCtx,
     const RunOptions& options) {
@@ -357,7 +357,7 @@ std::shared_ptr<runner::LocalRunner> SqlQueryRunner::makeLocalRunner(
       optimizerPool_);
 }
 
-std::vector<velox::RowVectorPtr> SqlQueryRunner::runSql(
+std::vector<velox::RowVectorPtr> AxiomSqlQueryRunner::runSql(
     const logical_plan::LogicalPlanNodePtr& logicalPlan,
     const RunOptions& options) {
   auto queryCtx = newQuery(options);

axiom/cli/SqlQueryRunner.h

@@ -26,21 +26,7 @@ namespace axiom::sql {
 
 class SqlQueryRunner {
  public:
-  /// @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.
-  void initialize(
-      const std::function<std::pair<std::string, std::optional<std::string>>(
-          facebook::axiom::optimizer::VeloxHistory& history)>&
-          initializeConnectors);
-
-  /// Results of running a query. SELECT queries return a vector of results.
-  /// Other queries return a message. SELECT query that returns no rows returns
-  /// std::nullopt message and empty vector of results.
-  struct SqlResult {
-    std::optional<std::string> message;
-    std::vector<facebook::velox::RowVectorPtr> results;
-  };
+  virtual ~SqlQueryRunner() = default;
 
   struct RunOptions {
     int32_t numWorkers{4};
@@ -52,30 +38,70 @@ class SqlQueryRunner {
     bool debugMode{false};
   };
 
+  /// Results of running a query. SELECT queries return a vector of results.
+  /// Other queries return a message. SELECT query that returns no rows returns
+  /// std::nullopt message and empty vector of results.
+  struct SqlResult {
+    std::optional<std::string> message;
+    std::vector<facebook::velox::RowVectorPtr> results;
+  };
+
+  /// Runs a single SQL statement and returns the result.
+  virtual SqlResult run(std::string_view sql, const RunOptions& options) = 0;
+
+  /// Runs a single parsed SQL statement and returns the result.
+  virtual SqlResult run(
+      const presto::SqlStatement& statement,
+      const RunOptions& options) = 0;
+
+  /// Parses SQL text containing one or more semicolon-separated statements.
+  /// @param sql SQL text to parse.
+  /// @return Vector of parsed statements.
+  virtual std::vector<presto::SqlStatementPtr> parseMultiple(
+      std::string_view sql,
+      const RunOptions& options) = 0;
+
+  virtual std::unordered_map<std::string, std::string>& sessionConfig() = 0;
+
+  virtual void saveHistory(const std::string& path) = 0;
+
+  virtual void clearHistory() = 0;
+};
+
+class AxiomSqlQueryRunner : public SqlQueryRunner {
+ public:
+  /// @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.
+  void initialize(
+      const std::function<std::pair<std::string, std::optional<std::string>>(
+          facebook::axiom::optimizer::VeloxHistory& history)>&
+          initializeConnectors);
+
   /// Runs a single SQL statement and returns the result.
-  SqlResult run(std::string_view sql, const RunOptions& options);
+  SqlResult run(std::string_view sql, const RunOptions& options) override;
 
   /// Runs a single parsed SQL statement and returns the result.
   SqlResult run(
       const presto::SqlStatement& statement,
-      const RunOptions& options);
+      const RunOptions& options) override;
 
   /// Parses SQL text containing one or more semicolon-separated statements.
   /// @param sql SQL text to parse.
   /// @return Vector of parsed statements.
   std::vector<presto::SqlStatementPtr> parseMultiple(
       std::string_view sql,
-      const RunOptions& options);
+      const RunOptions& options) override;
 
-  std::unordered_map<std::string, std::string>& sessionConfig() {
+  std::unordered_map<std::string, std::string>& sessionConfig() override {
     return config_;
   }
 
-  void saveHistory(const std::string& path) {
+  void saveHistory(const std::string& path) override {
     history_->saveToFile(path);
   }
 
-  void clearHistory() {
+  void clearHistory() override {
     history_ = std::make_unique<facebook::axiom::optimizer::VeloxHistory>();
   }
 

axiom/cli/tests/SqlQueryRunnerTest.cpp

@@ -38,7 +38,7 @@ class SqlQueryRunnerTest : public ::testing::Test, public test::VectorTestBase {
   }
 
   std::unique_ptr<SqlQueryRunner> makeRunner() {
-    auto runner = std::make_unique<SqlQueryRunner>();
+    auto runner = std::make_unique<AxiomSqlQueryRunner>();
 
     runner->initialize([&](auto&) {
       static int32_t kCounter = 0;