feat(parser): Add input/output table extraction to PrestoParser (#804)

Summary:

There are usecases for which callers may want to extract some information from a query without needing to resolve all the metadata details required to build a full logical plan. An example could be a client-side check decides where to send a query based on the tables it accesses, or moving ACL checks earlier in query execution by determining accessed tables immediately

See #789 for the related issue

To enable this, I am adding two APIs to the PrestoParser, one which extracts accessed input tables, and one which extracts output tables, if any exist

There are two parts to this changeset:

1. on recommendation of Masha, defined a DefaultTraversalVisitor, which performs a DFS traversal over all nodes in the AST. I used this baseclass for the existing ExprAnalyzer and the new TableVisitor. I can pull this into a separate PR if desired
2. add the TableVisitor, which extracts input tables and the output table for the query, and link it into two new PrestoParser APIs for input and output tables respectively

Some things I was unsure about and would like feedback:

1. I exposed two APIs, but I could easily have exposed one (getInputAndOutputTables) and return a struct containing the output of both APIs
2. I implemented the handlers for query types not currently covered by the parser (materialized view statements, some view statements, pure CREATE TABLE), but these cannot be run yet. I can also remove them or leave more comments in PrestoParser.cpp

I am also looking for comments on structuring: PrestoParser.cpp is getting big, I can cut it up into a few source/header files in this diff or a follow-up if others agree (but did not want to do so without discussion)

Reviewed By: mbasmanova

Differential Revision: D91525572

Author: hdikeman

axiom/sql/presto/CMakeLists.txt

@@ -16,7 +16,7 @@ add_subdirectory(grammar)
 add_subdirectory(ast)
 add_subdirectory(example)
 
-add_library(axiom_sql_presto_parser PrestoParser.cpp SqlStatement.cpp)
+add_library(axiom_sql_presto_parser PrestoParser.cpp SqlStatement.cpp TableVisitor.cpp)
 
 target_link_libraries(
   axiom_sql_presto_parser

axiom/sql/presto/PrestoParser.cpp

@@ -20,6 +20,7 @@
 #include "axiom/connectors/ConnectorMetadata.h"
 #include "axiom/logical_plan/PlanBuilder.h"
 #include "axiom/sql/presto/PrestoParseError.h"
+#include "axiom/sql/presto/TableVisitor.h"
 #include "axiom/sql/presto/ast/AstBuilder.h"
 #include "axiom/sql/presto/ast/AstPrinter.h"
 #include "axiom/sql/presto/ast/DefaultTraversalVisitor.h"
@@ -2201,4 +2202,19 @@ SqlStatementPtr PrestoParser::doParse(
   return doPlan(query, defaultConnectorId_, defaultSchema_, parseSql);
 }
 
+ReferencedTables PrestoParser::getReferencedTables(std::string_view sql) {
+  ParserHelper helper(sql);
+  auto* context = helper.parse();
+
+  AstBuilder astBuilder(false);
+  auto statement =
+      std::any_cast<std::shared_ptr<Statement>>(astBuilder.visit(context));
+
+  TableVisitor visitor(defaultConnectorId_, defaultSchema_);
+  visitor.process(statement.get());
+  return ReferencedTables{
+      .inputTables = visitor.inputTables(),
+      .outputTable = visitor.outputTable()};
+}
+
 } // namespace axiom::sql::presto

axiom/sql/presto/PrestoParser.h

@@ -19,6 +19,18 @@
 
 namespace axiom::sql::presto {
 
+/// The set of tables that a SQL statement references. Table names
+/// are returned in fully-qualified format ('catalog.schema.table').
+struct ReferencedTables {
+  /// The set of tables accessed for reading by the query,
+  /// or the empty set if the query does not read any tables.
+  std::unordered_set<std::string> inputTables;
+
+  /// Any table which would be modified by the query, or
+  /// nullopt if the query does not modify any tables.
+  std::optional<std::string> outputTable;
+};
+
 /// SQL Parser compatible with PrestoSQL dialect.
 class PrestoParser {
  public:
@@ -47,6 +59,13 @@ class PrestoParser {
       std::string_view sql,
       bool enableTracing = false);
 
+  /// Extracts tables referenced in a SQL statement, if any exist. This includes
+  /// table references which could later be optimized out, if their results
+  /// do not affect the query output (e.g., an unreferenced CTE).
+  /// @param sql SQL query statement
+  /// @return input and output tables which the query references.
+  ReferencedTables getReferencedTables(std::string_view sql);
+
  private:
   SqlStatementPtr doParse(std::string_view sql, bool enableTracing);
 

axiom/sql/presto/TableVisitor.cpp

@@ -0,0 +1,126 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "axiom/sql/presto/TableVisitor.h"
+
+#include <fmt/format.h>
+
+#include "velox/common/base/Exceptions.h"
+
+namespace axiom::sql::presto {
+
+TableVisitor::TableVisitor(
+    const std::string& defaultConnectorId,
+    const std::optional<std::string>& defaultSchema)
+    : defaultConnectorId_(defaultConnectorId), defaultSchema_(defaultSchema) {}
+
+void TableVisitor::visitWithQuery(WithQuery* node) {
+  // To cover the case where a CTE aliases an underlying
+  // table, e.g. 'WITH t AS (SELECT * FROM t)', we need to
+  // traverse the inner query before tracking the CTE alias.
+  DefaultTraversalVisitor::visitWithQuery(node);
+  ctes_.insert(node->name()->value());
+}
+
+void TableVisitor::visitTable(Table* node) {
+  const auto& parts = node->name()->parts();
+  if (parts.size() == 1 && ctes_.count(parts[0]) > 0) {
+    return;
+  }
+  inputTables_.insert(constructTableName(*node->name()));
+  DefaultTraversalVisitor::visitTable(node);
+}
+
+void TableVisitor::visitInsert(Insert* node) {
+  setOutputTable(*node->target());
+  DefaultTraversalVisitor::visitInsert(node);
+}
+
+void TableVisitor::visitCreateTableAsSelect(CreateTableAsSelect* node) {
+  setOutputTable(*node->name());
+  DefaultTraversalVisitor::visitCreateTableAsSelect(node);
+}
+
+void TableVisitor::visitUpdate(Update* node) {
+  setOutputTable(*node->table());
+  DefaultTraversalVisitor::visitUpdate(node);
+}
+
+void TableVisitor::visitDelete(Delete* node) {
+  setOutputTable(*node->table());
+  DefaultTraversalVisitor::visitDelete(node);
+}
+
+void TableVisitor::visitCreateTable(CreateTable* node) {
+  setOutputTable(*node->name());
+  DefaultTraversalVisitor::visitCreateTable(node);
+}
+
+void TableVisitor::visitCreateView(CreateView* node) {
+  setOutputTable(*node->name());
+  DefaultTraversalVisitor::visitCreateView(node);
+}
+
+void TableVisitor::visitCreateMaterializedView(CreateMaterializedView* node) {
+  setOutputTable(*node->name());
+  DefaultTraversalVisitor::visitCreateMaterializedView(node);
+}
+
+void TableVisitor::visitDropTable(DropTable* node) {
+  setOutputTable(*node->tableName());
+  DefaultTraversalVisitor::visitDropTable(node);
+}
+
+void TableVisitor::visitDropView(DropView* node) {
+  setOutputTable(*node->viewName());
+  DefaultTraversalVisitor::visitDropView(node);
+}
+
+void TableVisitor::visitDropMaterializedView(DropMaterializedView* node) {
+  setOutputTable(*node->viewName());
+  DefaultTraversalVisitor::visitDropMaterializedView(node);
+}
+
+std::string TableVisitor::constructTableName(const QualifiedName& name) const {
+  const auto& parts = name.parts();
+  VELOX_CHECK(!parts.empty(), "Table name cannot be empty");
+  VELOX_CHECK_LE(
+      parts.size(),
+      3,
+      "Table name must have 1-3 components, '{}'",
+      name.fullyQualifiedName());
+  switch (parts.size()) {
+    case 1:
+      if (defaultSchema_.has_value()) {
+        return fmt::format(
+            "{}.{}.{}", defaultConnectorId_, defaultSchema_.value(), parts[0]);
+      }
+      return fmt::format("{}.{}", defaultConnectorId_, parts[0]);
+    case 2:
+      return fmt::format("{}.{}.{}", defaultConnectorId_, parts[0], parts[1]);
+    case 3:
+      return fmt::format("{}.{}.{}", parts[0], parts[1], parts[2]);
+    default:
+      VELOX_UNREACHABLE();
+  }
+}
+
+void TableVisitor::setOutputTable(const QualifiedName& name) {
+  VELOX_CHECK(!outputTable_.has_value());
+  outputTable_ = constructTableName(name);
+}
+
+} // namespace axiom::sql::presto

axiom/sql/presto/TableVisitor.h

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <unordered_set>
+
+#include "axiom/sql/presto/ast/DefaultTraversalVisitor.h"
+
+namespace axiom::sql::presto {
+
+// Analyzes an expression to extract the fully-qualified names of any
+// input or output tables or views in the expression. Table accesses
+// inside CTEs are included, even if the CTE is never read from.
+class TableVisitor : public DefaultTraversalVisitor {
+ public:
+  TableVisitor(
+      const std::string& defaultConnectorId,
+      const std::optional<std::string>& defaultSchema);
+
+  const std::unordered_set<std::string>& inputTables() const {
+    return inputTables_;
+  }
+
+  const std::optional<std::string>& outputTable() const {
+    return outputTable_;
+  }
+
+ protected:
+  void visitWithQuery(WithQuery* node) override;
+  void visitTable(Table* node) override;
+  void visitInsert(Insert* node) override;
+  void visitCreateTableAsSelect(CreateTableAsSelect* node) override;
+  void visitUpdate(Update* node) override;
+  void visitDelete(Delete* node) override;
+  void visitCreateTable(CreateTable* node) override;
+  void visitCreateView(CreateView* node) override;
+  void visitCreateMaterializedView(CreateMaterializedView* node) override;
+  void visitDropTable(DropTable* node) override;
+  void visitDropView(DropView* node) override;
+  void visitDropMaterializedView(DropMaterializedView* node) override;
+
+ private:
+  std::string constructTableName(const QualifiedName& name) const;
+  void setOutputTable(const QualifiedName& name);
+
+  const std::string& defaultConnectorId_;
+  const std::optional<std::string>& defaultSchema_;
+  std::unordered_set<std::string> ctes_;
+  std::unordered_set<std::string> inputTables_;
+  std::optional<std::string> outputTable_;
+};
+
+} // namespace axiom::sql::presto

axiom/sql/presto/tests/CMakeLists.txt

@@ -12,7 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-add_executable(axiom_sql_presto_test LogicalPlanMatcher.cpp PrestoParserTest.cpp)
+add_executable(
+  axiom_sql_presto_test
+  LogicalPlanMatcher.cpp
+  PrestoParserTest.cpp
+  TableExtractorTest.cpp
+)
 
 add_test(axiom_sql_presto_test axiom_sql_presto_test)