tp: implement CREATE PERFETTO MACRO

This CL adds support for parsing and executing macros in PerfettoSQL
statements. This feature is important so that users can pass tables as
arguments which is not currently possible.

Bug: 290185551
Change-Id: I7519d85b5aadf602a167a6b45cb9e3a61ae0391a

src/trace_processor/perfetto_sql/engine/perfetto_sql_engine.cc

@@ -20,6 +20,7 @@
 #include <optional>
 #include <string>
 #include <variant>
+#include <vector>
 
 #include "perfetto/base/status.h"
 #include "perfetto/ext/base/status_or.h"
@@ -220,7 +221,7 @@ PerfettoSqlEngine::ExecuteUntilLastStatement(SqlSource sql_source) {
   //    statement for the last valid statement.
   std::optional<SqliteEngine::PreparedStatement> res;
   ExecutionStats stats;
-  PerfettoSqlParser parser(std::move(sql_source));
+  PerfettoSqlParser parser(std::move(sql_source), macros_);
   while (parser.Next()) {
     std::optional<SqlSource> source;
     if (auto* cf = std::get_if<PerfettoSqlParser::CreateFunction>(
@@ -238,6 +239,11 @@ PerfettoSqlEngine::ExecuteUntilLastStatement(SqlSource sql_source) {
                    &parser.statement())) {
       RETURN_IF_ERROR(ExecuteInclude(*include, parser));
       source = RewriteToDummySql(parser.statement_sql());
+    } else if (auto* macro = std::get_if<PerfettoSqlParser::CreateMacro>(
+                   &parser.statement())) {
+      auto sql = macro->sql;
+      RETURN_IF_ERROR(ExecuteCreateMacro(*macro));
+      source = RewriteToDummySql(sql);
     } else {
       // If none of the above matched, this must just be an SQL statement
       // directly executable by SQLite.
@@ -597,6 +603,55 @@ base::StatusOr<SqlSource> PerfettoSqlEngine::ExecuteCreateFunction(
       SqlSource::FromTraceProcessorImplementation(create.ToStdString()));
 }
 
+base::Status PerfettoSqlEngine::ExecuteCreateMacro(
+    const PerfettoSqlParser::CreateMacro& create_macro) {
+  // Check that the argument types is one of the allowed types.
+  for (const auto& [name, type] : create_macro.args) {
+    std::string lower_type = base::ToLower(type.sql());
+    if (lower_type != "tableorsubquery" && lower_type != "expr") {
+      // TODO(lalitm): add a link to create macro documentation.
+      return base::ErrStatus(
+          "%sMacro %s argument %s is unkown type %s. Allowed types: "
+          "TableOrSubquery, Expr",
+          type.AsTraceback(0).c_str(), create_macro.name.sql().c_str(),
+          name.sql().c_str(), type.sql().c_str());
+    }
+  }
+  std::string lower_return = base::ToLower(create_macro.returns.sql());
+  if (lower_return != "tableorsubquery" && lower_return != "expr") {
+    // TODO(lalitm): add a link to create macro documentation.
+    return base::ErrStatus(
+        "%sMacro %s return type %s is unknown. Allowed types: "
+        "TableOrSubquery, Expr",
+        create_macro.returns.AsTraceback(0).c_str(),
+        create_macro.name.sql().c_str(), create_macro.returns.sql().c_str());
+  }
+
+  std::vector<std::string> args;
+  for (const auto& arg : create_macro.args) {
+    args.push_back(arg.first.sql());
+  }
+  PerfettoSqlPreprocessor::Macro macro{
+      create_macro.replace,
+      create_macro.name.sql(),
+      std::move(args),
+      create_macro.sql,
+  };
+  if (auto it = macros_.Find(create_macro.name.sql()); it) {
+    if (!create_macro.replace) {
+      // TODO(lalitm): add a link to create macro documentation.
+      return base::ErrStatus("%sMacro already exists",
+                             create_macro.name.AsTraceback(0).c_str());
+    }
+    *it = std::move(macro);
+    return base::OkStatus();
+  }
+  std::string name = macro.name;
+  auto it_and_inserted = macros_.Insert(std::move(name), std::move(macro));
+  PERFETTO_CHECK(it_and_inserted.second);
+  return base::OkStatus();
+}
+
 RuntimeTableFunction::State* PerfettoSqlEngine::GetRuntimeTableFunctionState(
     const std::string& name) const {
   auto it = runtime_table_fn_states_.Find(base::ToLower(name));

src/trace_processor/perfetto_sql/engine/perfetto_sql_engine.h

@@ -152,12 +152,15 @@ class PerfettoSqlEngine {
   // Registers a SQL-defined trace processor C++ table with SQLite.
   base::Status RegisterRuntimeTable(std::string name, SqlSource sql);
 
+  base::Status ExecuteCreateMacro(const PerfettoSqlParser::CreateMacro&);
+
   std::unique_ptr<QueryCache> query_cache_;
   StringPool* pool_ = nullptr;
   base::FlatHashMap<std::string, std::unique_ptr<RuntimeTableFunction::State>>
       runtime_table_fn_states_;
   base::FlatHashMap<std::string, std::unique_ptr<RuntimeTable>> runtime_tables_;
   base::FlatHashMap<std::string, sql_modules::RegisteredModule> modules_;
+  base::FlatHashMap<std::string, PerfettoSqlPreprocessor::Macro> macros_;
   std::unique_ptr<SqliteEngine> engine_;
 };
 

src/trace_processor/perfetto_sql/engine/perfetto_sql_engine_unittest.cc

@@ -109,6 +109,24 @@ TEST_F(PerfettoSqlEngineTest, CreateTableFunctionDupe) {
   ASSERT_TRUE(res.ok());
 }
 
+TEST_F(PerfettoSqlEngineTest, CreateMacro) {
+  auto res_create = engine_.Execute(SqlSource::FromExecuteQuery(
+      "CREATE PERFETTO MACRO foo() RETURNS TableOrSubquery AS select 42 AS x"));
+  ASSERT_TRUE(res_create.ok()) << res_create.status().c_message();
+
+  res_create = engine_.Execute(SqlSource::FromExecuteQuery(
+      "CREATE PERFETTO MACRO bar(x TableOrSubquery) RETURNS TableOrSubquery AS "
+      "select * from $x"));
+  ASSERT_TRUE(res_create.ok()) << res_create.status().c_message();
+
+  auto res = engine_.ExecuteUntilLastStatement(
+      SqlSource::FromExecuteQuery("bar!((foo!()))"));
+  ASSERT_TRUE(res.ok()) << res.status().c_message();
+  ASSERT_FALSE(res->stmt.IsDone());
+  ASSERT_EQ(sqlite3_column_int64(res->stmt.sqlite_stmt(), 0), 42);
+  ASSERT_FALSE(res->stmt.Step());
+}
+
 }  // namespace
 }  // namespace trace_processor
 }  // namespace perfetto

src/trace_processor/perfetto_sql/engine/perfetto_sql_parser.cc

@@ -19,9 +19,14 @@
 #include <algorithm>
 #include <functional>
 #include <optional>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
 
 #include "perfetto/base/logging.h"
 #include "perfetto/base/status.h"
+#include "perfetto/ext/base/flat_hash_map.h"
 #include "perfetto/ext/base/status_or.h"
 #include "perfetto/ext/base/string_utils.h"
 #include "src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.h"
@@ -79,10 +84,28 @@ bool ValidateModuleName(const std::string& name) {
                       std::not_fn(IsValidModuleWord)) == packages.end();
 }
 
+std::string SerializeArgs(std::vector<std::pair<SqlSource, SqlSource>> args) {
+  bool comma = false;
+  std::string serialized;
+  for (const auto& [name, type] : args) {
+    if (comma) {
+      serialized.append(", ");
+    }
+    comma = true;
+    serialized.append(name.sql().c_str());
+    serialized.push_back(' ');
+    serialized.append(type.sql().c_str());
+  }
+  return serialized;
+}
+
 }  // namespace
 
-PerfettoSqlParser::PerfettoSqlParser(SqlSource source)
-    : preprocessor_(std::move(source)),
+PerfettoSqlParser::PerfettoSqlParser(
+    SqlSource source,
+    const base::FlatHashMap<std::string, PerfettoSqlPreprocessor::Macro>&
+        macros)
+    : preprocessor_(std::move(source), macros),
       tokenizer_(SqlSource::FromTraceProcessorImplementation("")) {}
 
 bool PerfettoSqlParser::Next() {
@@ -188,9 +211,13 @@ bool PerfettoSqlParser::Next() {
         if (TokenIsSqliteKeyword("table", token)) {
           return ParseCreatePerfettoTable(*first_non_space_token);
         }
+        if (TokenIsCustomKeyword("macro", token)) {
+          return ParseCreatePerfettoMacro(state ==
+                                          State::kCreateOrReplacePerfetto);
+        }
         base::StackString<1024> err(
-            "Expected 'FUNCTION' or 'TABLE' after 'CREATE PERFETTO', received "
-            "'%*s'.",
+            "Expected 'FUNCTION', 'TABLE' or 'MACRO' after 'CREATE PERFETTO', "
+            "received '%*s'.",
             static_cast<int>(token.str.size()), token.str.data());
         return ErrorAtToken(token, err.c_str());
     }
@@ -263,10 +290,13 @@ bool PerfettoSqlParser::ParseCreatePerfettoFunction(
     return ErrorAtToken(lp, "Malformed function prototype: '(' expected");
   }
 
-  prototype.push_back('(');
-  if (!ParseArgumentDefinitions(&prototype)) {
+  std::vector<Argument> args;
+  if (!ParseArgumentDefinitions(args)) {
     return false;
   }
+
+  prototype.push_back('(');
+  prototype.append(SerializeArgs(args));
   prototype.push_back(')');
 
   if (Token returns = tokenizer_.NextNonWhitespace();
@@ -285,9 +315,11 @@ bool PerfettoSqlParser::ParseCreatePerfettoFunction(
       return ErrorAtToken(lp, "Malformed table return: '(' expected");
     }
     // Table function return.
-    if (!ParseArgumentDefinitions(&ret)) {
+    std::vector<Argument> ret_args;
+    if (!ParseArgumentDefinitions(ret_args)) {
       return false;
     }
+    ret = SerializeArgs(ret_args);
   } else if (ret_token.token_type != SqliteTokenType::TK_ID) {
     // TODO(lalitm): add a link to create function documentation.
     return ErrorAtToken(ret_token, "Invalid return type");
@@ -310,28 +342,117 @@ bool PerfettoSqlParser::ParseCreatePerfettoFunction(
   return true;
 }
 
-bool PerfettoSqlParser::ParseArgumentDefinitions(std::string* str) {
-  for (Token tok = tokenizer_.Next();; tok = tokenizer_.Next()) {
-    if (tok.token_type == SqliteTokenType::TK_RP) {
-      return true;
+bool PerfettoSqlParser::ParseCreatePerfettoMacro(bool replace) {
+  Token name = tokenizer_.NextNonWhitespace();
+  if (name.token_type != SqliteTokenType::TK_ID) {
+    // TODO(lalitm): add a link to create macro documentation.
+    base::StackString<1024> err("Invalid macro name %.*s",
+                                static_cast<int>(name.str.size()),
+                                name.str.data());
+    return ErrorAtToken(name, err.c_str());
+  }
+
+  // TK_LP == '(' (i.e. left parenthesis).
+  if (Token lp = tokenizer_.NextNonWhitespace();
+      lp.token_type != SqliteTokenType::TK_LP) {
+    // TODO(lalitm): add a link to create macro documentation.
+    return ErrorAtToken(lp, "Malformed macro prototype: '(' expected");
+  }
+
+  std::vector<Argument> args;
+  if (!ParseArgumentDefinitions(args)) {
+    return false;
+  }
+
+  if (Token returns = tokenizer_.NextNonWhitespace();
+      !TokenIsCustomKeyword("returns", returns)) {
+    // TODO(lalitm): add a link to create macro documentation.
+    return ErrorAtToken(returns, "Expected keyword 'returns'");
+  }
+
+  Token returns_value = tokenizer_.NextNonWhitespace();
+  if (returns_value.token_type != SqliteTokenType::TK_ID) {
+    // TODO(lalitm): add a link to create function documentation.
+    return ErrorAtToken(returns_value, "Expected return type");
+  }
+
+  if (Token as_token = tokenizer_.NextNonWhitespace();
+      !TokenIsSqliteKeyword("as", as_token)) {
+    // TODO(lalitm): add a link to create macro documentation.
+    return ErrorAtToken(as_token, "Expected keyword 'as'");
+  }
+
+  Token first = tokenizer_.NextNonWhitespace();
+  Token tok = tokenizer_.NextTerminal();
+  statement_ = CreateMacro{
+      replace, tokenizer_.SubstrToken(name), std::move(args),
+      tokenizer_.SubstrToken(returns_value), tokenizer_.Substr(first, tok)};
+  return true;
+}
+
+bool PerfettoSqlParser::ParseArgumentDefinitions(std::vector<Argument>& res) {
+  enum TokenType {
+    kIdOrRp,
+    kId,
+    kType,
+    kCommaOrRp,
+  };
+
+  std::optional<Token> id = std::nullopt;
+  TokenType expected = kIdOrRp;
+  for (Token tok = tokenizer_.NextNonWhitespace();;
+       tok = tokenizer_.NextNonWhitespace()) {
+    // Keywords can be used as names accidentally so have an explicit error
+    // message for those.
+    if (tok.token_type == SqliteTokenType::TK_GENERIC_KEYWORD) {
+      base::StackString<1024> err(
+          "Malformed function/macro prototype: %.*s is a SQL keyword so "
+          "cannot appear in a prototype",
+          static_cast<int>(tok.str.size()), tok.str.data());
+      return ErrorAtToken(tok, err.c_str());
     }
-    if (tok.token_type == SqliteTokenType::TK_SPACE) {
-      str->append(" ");
-      continue;
+    if (expected == kCommaOrRp) {
+      PERFETTO_CHECK(expected == kCommaOrRp);
+      if (tok.token_type == SqliteTokenType::TK_RP) {
+        return true;
+      }
+      if (tok.token_type == SqliteTokenType::TK_COMMA) {
+        expected = kId;
+        continue;
+      }
+      return ErrorAtToken(tok, "')' or ',' expected");
     }
-    if (tok.token_type != SqliteTokenType::TK_ID &&
-        tok.token_type != SqliteTokenType::TK_COMMA) {
-      if (tok.token_type == SqliteTokenType::TK_GENERIC_KEYWORD) {
-        base::StackString<1024> err(
-            "Malformed function prototype: %.*s is a SQL keyword so cannot "
-            "appear in a function prototype",
-            static_cast<int>(tok.str.size()), tok.str.data());
+    if (expected == kType) {
+      if (tok.token_type != SqliteTokenType::TK_ID) {
+        // TODO(lalitm): add a link to documentation.
+        base::StackString<1024> err("%.*s is not a valid argument type",
+                                    static_cast<int>(tok.str.size()),
+                                    tok.str.data());
         return ErrorAtToken(tok, err.c_str());
       }
-      // TODO(lalitm): add a link to create function documentation.
-      return ErrorAtToken(tok, "')', ',', name or type expected");
+      PERFETTO_CHECK(id);
+      res.push_back(std::make_pair(tokenizer_.SubstrToken(*id),
+                                   tokenizer_.SubstrToken(tok)));
+      id = std::nullopt;
+      expected = kCommaOrRp;
+      continue;
+    }
+
+    // kIdOrRp only happens on the very first token.
+    if (tok.token_type == SqliteTokenType::TK_RP && expected == kIdOrRp) {
+      return true;
+    }
+
+    if (tok.token_type != SqliteTokenType::TK_ID) {
+      // TODO(lalitm): add a link to documentation.
+      base::StackString<1024> err("%.*s is not a valid argument name",
+                                  static_cast<int>(tok.str.size()),
+                                  tok.str.data());
+      return ErrorAtToken(tok, err.c_str());
     }
-    str->append(tok.str);
+    id = tok;
+    expected = kType;
+    continue;
   }
 }