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

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)

Differential Revision: D91525572

Author: hdikeman

axiom/sql/presto/PrestoParser.cpp

@@ -22,6 +22,7 @@
 #include "axiom/sql/presto/PrestoParseError.h"
 #include "axiom/sql/presto/ast/AstBuilder.h"
 #include "axiom/sql/presto/ast/AstPrinter.h"
+#include "axiom/sql/presto/ast/DefaultTraversalVisitor.h"
 #include "axiom/sql/presto/ast/UpperCaseInputStream.h"
 #include "axiom/sql/presto/grammar/PrestoSqlLexer.h"
 #include "axiom/sql/presto/grammar/PrestoSqlParser.h"
@@ -179,59 +180,20 @@ std::string canonicalizeIdentifier(const Identifier& identifier) {
   return canonicalizeName(identifier.value());
 }
 
-// Analizes the expression to find out whether there are any aggregate function
+// Analyzes the expression to find out whether there are any aggregate function
 // calls and to verify that aggregate calls are not nested, e.g. sum(count(x))
 // is not allowed.
-class ExprAnalyzer : public AstVisitor {
+class ExprAnalyzer : public DefaultTraversalVisitor {
  public:
   bool hasAggregate() const {
     return numAggregates_ > 0;
   }
 
- private:
-  void defaultVisit(Node* node) override {
-    if (dynamic_cast<Literal*>(node) != nullptr) {
-      // Literals have no function calls.
-      return;
-    }
-
-    VELOX_NYI(
-        "Not yet supported node type: {}", NodeTypeName::toName(node->type()));
-  }
-
-  void visitArithmeticUnaryExpression(
-      ArithmeticUnaryExpression* node) override {
-    node->value()->accept(this);
-  }
-
-  void visitArrayConstructor(ArrayConstructor* node) override {
-    for (const auto& value : node->values()) {
-      value->accept(this);
-    }
-  }
-
-  void visitBetweenPredicate(BetweenPredicate* node) override {
-    node->value()->accept(this);
-    node->min()->accept(this);
-    node->max()->accept(this);
-  }
-
-  void visitCast(Cast* node) override {
-    node->expression()->accept(this);
-  }
-
-  void visitDereferenceExpression(DereferenceExpression* node) override {
-    node->base()->accept(this);
-  }
-
+ protected:
   void visitExistsPredicate(ExistsPredicate* node) override {
     // Aggregate function calls within a subquery do not count.
   }
 
-  void visitExtract(Extract* node) override {
-    node->expression()->accept(this);
-  }
-
   void visitFunctionCall(FunctionCall* node) override {
     const auto& name = node->name()->suffix();
     if (facebook::velox::exec::getAggregateFunctionEntry(name)) {
@@ -245,112 +207,16 @@ class ExprAnalyzer : public AstVisitor {
       ++numAggregates_;
     }
 
-    for (const auto& arg : node->arguments()) {
-      arg->accept(this);
-    }
+    DefaultTraversalVisitor::visitFunctionCall(node);
 
     aggregateName_.reset();
   }
 
-  void visitInListExpression(InListExpression* node) override {
-    for (const auto& value : node->values()) {
-      value->accept(this);
-    }
-  }
-
-  void visitInPredicate(InPredicate* node) override {
-    node->value()->accept(this);
-    node->valueList()->accept(this);
-  }
-
-  void visitIsNullPredicate(IsNullPredicate* node) override {
-    node->value()->accept(this);
-  }
-
-  void visitIsNotNullPredicate(IsNotNullPredicate* node) override {
-    node->value()->accept(this);
-  }
-
-  void visitLambdaExpression(LambdaExpression* node) override {
-    node->body()->accept(this);
-  }
-
-  void visitNotExpression(NotExpression* node) override {
-    node->value()->accept(this);
-  }
-
-  void visitArithmeticBinaryExpression(
-      ArithmeticBinaryExpression* node) override {
-    node->left()->accept(this);
-    node->right()->accept(this);
-  }
-
-  void visitLogicalBinaryExpression(LogicalBinaryExpression* node) override {
-    node->left()->accept(this);
-    node->right()->accept(this);
-  }
-
-  void visitComparisonExpression(ComparisonExpression* node) override {
-    node->left()->accept(this);
-    node->right()->accept(this);
-  }
-
-  void visitLikePredicate(LikePredicate* node) override {
-    node->value()->accept(this);
-    node->pattern()->accept(this);
-    if (node->escape() != nullptr) {
-      node->escape()->accept(this);
-    }
-  }
-
-  void visitSimpleCaseExpression(SimpleCaseExpression* node) override {
-    node->operand()->accept(this);
-
-    for (const auto& clause : node->whenClauses()) {
-      clause->operand()->accept(this);
-      clause->result()->accept(this);
-    }
-
-    if (node->defaultValue()) {
-      node->defaultValue()->accept(this);
-    }
-  }
-
-  void visitSearchedCaseExpression(SearchedCaseExpression* node) override {
-    for (const auto& clause : node->whenClauses()) {
-      clause->operand()->accept(this);
-      clause->result()->accept(this);
-    }
-
-    if (node->defaultValue()) {
-      node->defaultValue()->accept(this);
-    }
-  }
-
   void visitSubqueryExpression(SubqueryExpression* node) override {
     // Aggregate function calls within a subquery do not count.
   }
 
-  void visitSubscriptExpression(SubscriptExpression* node) override {
-    node->base()->accept(this);
-    node->index()->accept(this);
-  }
-
-  void visitIdentifier(Identifier* node) override {
-    // No function calls.
-  }
-
-  void visitRow(Row* node) override {
-    for (const auto& item : node->items()) {
-      item->accept(this);
-    }
-  }
-
-  void visitAtTimeZone(AtTimeZone* node) override {
-    node->value()->accept(this);
-    node->timeZone()->accept(this);
-  }
-
+ private:
   size_t numAggregates_{0};
   std::optional<std::string> aggregateName_;
 };
@@ -2335,4 +2201,169 @@ SqlStatementPtr PrestoParser::doParse(
   return doPlan(query, defaultConnectorId_, defaultSchema_, parseSql);
 }
 
+namespace {
+
+// 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)
+      : defaultConnectorId_(defaultConnectorId),
+        defaultSchema_(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 {
+    // 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 visitTable(Table* node) override {
+    if (isCteReference(*node->name())) {
+      return;
+    }
+    inputTables_.insert(constructTableName(*node->name()));
+    DefaultTraversalVisitor::visitTable(node);
+  }
+
+  void visitInsert(Insert* node) override {
+    parseOutputTable(*node->target());
+    DefaultTraversalVisitor::visitInsert(node);
+  }
+
+  void visitCreateTableAsSelect(CreateTableAsSelect* node) override {
+    parseOutputTable(*node->name());
+    DefaultTraversalVisitor::visitCreateTableAsSelect(node);
+  }
+
+  void visitUpdate(Update* node) override {
+    parseOutputTable(*node->table());
+    DefaultTraversalVisitor::visitUpdate(node);
+  }
+
+  void visitDelete(Delete* node) override {
+    parseOutputTable(*node->table());
+    DefaultTraversalVisitor::visitDelete(node);
+  }
+
+  void visitCreateTable(CreateTable* node) override {
+    parseOutputTable(*node->name());
+    DefaultTraversalVisitor::visitCreateTable(node);
+  }
+
+  void visitCreateView(CreateView* node) override {
+    parseOutputTable(*node->name());
+    DefaultTraversalVisitor::visitCreateView(node);
+  }
+
+  void visitCreateMaterializedView(CreateMaterializedView* node) override {
+    parseOutputTable(*node->name());
+    DefaultTraversalVisitor::visitCreateMaterializedView(node);
+  }
+
+  void visitDropTable(DropTable* node) override {
+    parseOutputTable(*node->tableName());
+    DefaultTraversalVisitor::visitDropTable(node);
+  }
+
+  void visitDropView(DropView* node) override {
+    parseOutputTable(*node->viewName());
+    DefaultTraversalVisitor::visitDropView(node);
+  }
+
+  void visitDropMaterializedView(DropMaterializedView* node) override {
+    parseOutputTable(*node->viewName());
+    DefaultTraversalVisitor::visitDropMaterializedView(node);
+  }
+
+ private:
+  std::string constructTableName(const QualifiedName& name) {
+    const auto& parts = name.parts();
+    VELOX_CHECK(!parts.empty(), "Table name cannot be empty");
+    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_FAIL(
+            "Table name must have 1-3 components, '{}' has {}",
+            name.fullyQualifiedName(),
+            parts.size());
+    }
+  }
+
+  void parseOutputTable(const QualifiedName& name) {
+    auto outputTable = constructTableName(name);
+    VELOX_CHECK(
+        !outputTable_.has_value(),
+        "Cannot perform write against multiple tables (found {} and {})",
+        outputTable_.value(),
+        outputTable);
+    outputTable_ = std::move(outputTable);
+  }
+
+  bool isCteReference(const QualifiedName& name) {
+    const auto& parts = name.parts();
+    return parts.size() == 1 && ctes_.count(parts[0]) > 0;
+  }
+
+  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
+
+std::unordered_set<std::string> PrestoParser::getInputTables(
+    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 visitor.inputTables();
+}
+
+std::optional<std::string> PrestoParser::getOutputTable(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 visitor.outputTable();
+}
+
 } // namespace axiom::sql::presto

axiom/sql/presto/PrestoParser.h

@@ -47,6 +47,18 @@ class PrestoParser {
       std::string_view sql,
       bool enableTracing = false);
 
+  /// Extracts input tables from a SQL query, if any exist.
+  /// @param sql SQL query statement
+  /// @return set of fully-qualified input table names in `catalog.schema.table`
+  /// format. If the query accesses no tables, an empty set is returned
+  std::unordered_set<std::string> getInputTables(std::string_view sql);
+
+  /// Extracts the output table from a SQL query, if one exists.
+  /// @param sql SQL query statement.
+  /// @return fully-qualified output table name in `catalog.schema.table`
+  /// format, or std::nullopt if the query does not modify a table.
+  std::optional<std::string> getOutputTable(std::string_view sql);
+
  private:
   SqlStatementPtr doParse(std::string_view sql, bool enableTracing);