tp: Implement INCLUDE PERFETTO MODULE

Bug:296388363
Change-Id: I6fc43ceb4fd8d4e80b59c34409574e3f3af95824

src/trace_processor/perfetto_sql/engine/BUILD.gn

@@ -41,6 +41,7 @@ source_set("engine") {
     "../../types",
     "../../util",
     "../../util:sql_argument",
+    "../../util:stdlib",
   ]
 }
 

src/trace_processor/perfetto_sql/engine/perfetto_sql_engine.cc

@@ -192,24 +192,32 @@ PerfettoSqlEngine::ExecuteUntilLastStatement(SqlSource sql_source) {
     std::optional<SqlSource> source;
     if (auto* cf = std::get_if<PerfettoSqlParser::CreateFunction>(
             &parser.statement())) {
-      auto source_or = ExecuteCreateFunction(*cf);
-      RETURN_IF_ERROR(AddTracebackIfNeeded(source_or.status(), cf->sql));
+      auto source_or = ExecuteCreateFunction(*cf, parser);
+      RETURN_IF_ERROR(
+          AddTracebackIfNeeded(source_or.status(), parser.statement_sql()));
       source = std::move(source_or.value());
     } else if (auto* cst = std::get_if<PerfettoSqlParser::CreateTable>(
                    &parser.statement())) {
       RETURN_IF_ERROR(AddTracebackIfNeeded(
-          RegisterRuntimeTable(cst->name, cst->sql), cst->sql));
+          RegisterRuntimeTable(cst->name, cst->sql), parser.statement_sql()));
       // Since the rest of the code requires a statement, just use a no-value
       // dummy statement.
-      source = cst->sql.FullRewrite(
+      source = parser.statement_sql().FullRewrite(
+          SqlSource::FromTraceProcessorImplementation("SELECT 0 WHERE 0"));
+    } else if (auto* include = std::get_if<PerfettoSqlParser::Include>(
+                   &parser.statement())) {
+      RETURN_IF_ERROR(ExecuteInclude(*include, parser));
+      // Since the rest of the code requires a statement, just use a no-value
+      // dummy statement.
+      source = parser.statement_sql().FullRewrite(
           SqlSource::FromTraceProcessorImplementation("SELECT 0 WHERE 0"));
     } else {
       // If none of the above matched, this must just be an SQL statement
       // directly executable by SQLite.
       auto* sql =
           std::get_if<PerfettoSqlParser::SqliteSql>(&parser.statement());
       PERFETTO_CHECK(sql);
-      source = std::move(sql->sql);
+      source = parser.statement_sql();
     }
 
     // Try to get SQLite to prepare the statement.
@@ -399,15 +407,50 @@ base::Status PerfettoSqlEngine::EnableSqlFunctionMemoization(
   return CreatedFunction::EnableMemoization(ctx);
 }
 
+base::Status PerfettoSqlEngine::ExecuteInclude(
+    const PerfettoSqlParser::Include& include,
+    const PerfettoSqlParser& parser) {
+  std::string key = include.key;
+  PERFETTO_TP_TRACE(metatrace::Category::TOPLEVEL, "Import",
+                    [key](metatrace::Record* r) { r->AddArg("Import", key); });
+  std::string module_name = sql_modules::GetModuleName(key);
+  auto module = FindModule(module_name);
+  if (!module)
+    return base::ErrStatus("INCLUDE: Unknown module name provided - %s",
+                           key.c_str());
+
+  auto module_file = module->include_key_to_file.Find(key);
+  if (!module_file) {
+    return base::ErrStatus("INCLUDE: Unknown filename provided - %s",
+                           key.c_str());
+  }
+  // INCLUDE is noop for already included files.
+  if (module_file->included) {
+    return base::OkStatus();
+  }
+
+  auto it = Execute(SqlSource::FromModuleInclude(module_file->sql, key));
+  if (!it.status().ok()) {
+    return base::ErrStatus("%s%s",
+                           parser.statement_sql().AsTraceback(0).c_str(),
+                           it.status().c_message());
+  }
+  if (it->statement_count_with_output > 0)
+    return base::ErrStatus("INCLUDE: Included module returning values.");
+  module_file->included = true;
+  return base::OkStatus();
+}
+
 base::StatusOr<SqlSource> PerfettoSqlEngine::ExecuteCreateFunction(
-    const PerfettoSqlParser::CreateFunction& cf) {
+    const PerfettoSqlParser::CreateFunction& cf,
+    const PerfettoSqlParser& parser) {
   if (!cf.is_table) {
     RETURN_IF_ERROR(
         RegisterSqlFunction(cf.replace, cf.prototype, cf.returns, cf.sql));
 
     // Since the rest of the code requires a statement, just use a no-value
     // dummy statement.
-    return cf.sql.FullRewrite(
+    return parser.statement_sql().FullRewrite(
         SqlSource::FromTraceProcessorImplementation("SELECT 0 WHERE 0"));
   }
 

src/trace_processor/perfetto_sql/engine/perfetto_sql_engine.h

@@ -32,6 +32,7 @@
 #include "src/trace_processor/sqlite/sql_source.h"
 #include "src/trace_processor/sqlite/sqlite_engine.h"
 #include "src/trace_processor/sqlite/sqlite_utils.h"
+#include "src/trace_processor/util/sql_modules.h"
 
 namespace perfetto {
 namespace trace_processor {
@@ -128,9 +129,24 @@ class PerfettoSqlEngine {
 
   SqliteEngine* sqlite_engine() { return engine_.get(); }
 
+  // Makes new SQL module available to import.
+  void RegisterModule(const std::string& name,
+                      sql_modules::RegisteredModule module) {
+    modules_.Insert(name, std::move(module));
+  }
+
+  // Fetches registered SQL module.
+  sql_modules::RegisteredModule* FindModule(const std::string& name) {
+    return modules_.Find(name);
+  }
+
  private:
   base::StatusOr<SqlSource> ExecuteCreateFunction(
-      const PerfettoSqlParser::CreateFunction&);
+      const PerfettoSqlParser::CreateFunction&,
+      const PerfettoSqlParser& parser);
+
+  base::Status ExecuteInclude(const PerfettoSqlParser::Include&,
+                              const PerfettoSqlParser& parser);
 
   // Registers a SQL-defined trace processor C++ table with SQLite.
   base::Status RegisterRuntimeTable(std::string name, SqlSource sql);
@@ -140,6 +156,7 @@ class PerfettoSqlEngine {
   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_;
   std::unique_ptr<SqliteEngine> engine_;
 };
 

src/trace_processor/perfetto_sql/engine/perfetto_sql_parser.cc

@@ -17,6 +17,7 @@
 #include "src/trace_processor/perfetto_sql/engine/perfetto_sql_parser.h"
 
 #include <algorithm>
+#include <functional>
 #include <optional>
 
 #include "perfetto/base/logging.h"
@@ -37,6 +38,8 @@ using Statement = PerfettoSqlParser::Statement;
 enum class State {
   kStmtStart,
   kCreate,
+  kInclude,
+  kIncludePerfetto,
   kCreateOr,
   kCreateOrReplace,
   kCreateOrReplacePerfetto,
@@ -60,6 +63,21 @@ bool TokenIsCustomKeyword(std::string_view keyword, SqliteTokenizer::Token t) {
   return t.token_type == SqliteTokenType::TK_ID && KeywordEqual(keyword, t.str);
 }
 
+bool IsValidModuleWord(const std::string& word) {
+  for (const char& c : word) {
+    if (!std::isalnum(c) & (c != '_') & !std::islower(c)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool ValidateModuleName(const std::string& name) {
+  std::vector<std::string> packages = base::SplitString(name, ".");
+  return std::find_if(packages.begin(), packages.end(),
+                      std::not_fn(IsValidModuleWord)) == packages.end();
+}
+
 }  // namespace
 
 PerfettoSqlParser::PerfettoSqlParser(SqlSource sql)
@@ -81,8 +99,8 @@ bool PerfettoSqlParser::Next() {
       // If we have a non-space character we've seen, just return all the stuff
       // after that point.
       if (first_non_space_token) {
-        statement_ =
-            SqliteSql{tokenizer_.Substr(*first_non_space_token, token)};
+        statement_ = SqliteSql{};
+        statement_sql_ = tokenizer_.Substr(*first_non_space_token, token);
         return true;
       }
       // This means we've seen a semi-colon without any non-space content. Just
@@ -104,9 +122,29 @@ bool PerfettoSqlParser::Next() {
       case State::kPassthrough:
         break;
       case State::kStmtStart:
-        state = TokenIsSqliteKeyword("create", token) ? State::kCreate
-                                                      : State::kPassthrough;
+        if (TokenIsSqliteKeyword("create", token)) {
+          state = State::kCreate;
+        } else if (TokenIsCustomKeyword("include", token)) {
+          state = State::kInclude;
+        } else {
+          state = State::kPassthrough;
+        }
+        break;
+      case State::kInclude:
+        if (TokenIsCustomKeyword("perfetto", token)) {
+          state = State::kIncludePerfetto;
+        } else {
+          return ErrorAtToken(token,
+                              "Use 'INCLUDE PERFETTO MODULE {include_key}'.");
+        }
         break;
+      case State::kIncludePerfetto:
+        if (TokenIsCustomKeyword("module", token)) {
+          return ParseIncludePerfettoModule(*first_non_space_token);
+        } else {
+          return ErrorAtToken(token,
+                              "Use 'INCLUDE PERFETTO MODULE {include_key}'.");
+        }
       case State::kCreate:
         if (TokenIsSqliteKeyword("trigger", token)) {
           // TODO(lalitm): add this to the "errors" documentation page
@@ -135,11 +173,11 @@ bool PerfettoSqlParser::Next() {
       case State::kCreateOrReplacePerfetto:
       case State::kCreatePerfetto:
         if (TokenIsCustomKeyword("function", token)) {
-          return ParseCreatePerfettoFunction(state ==
-                                             State::kCreateOrReplacePerfetto);
+          return ParseCreatePerfettoFunction(
+              state == State::kCreateOrReplacePerfetto, *first_non_space_token);
         }
         if (TokenIsSqliteKeyword("table", token)) {
-          return ParseCreatePerfettoTable();
+          return ParseCreatePerfettoTable(*first_non_space_token);
         }
         base::StackString<1024> err(
             "Expected 'FUNCTION' or 'TABLE' after 'CREATE PERFETTO', received "
@@ -150,7 +188,26 @@ bool PerfettoSqlParser::Next() {
   }
 }
 
-bool PerfettoSqlParser::ParseCreatePerfettoTable() {
+bool PerfettoSqlParser::ParseIncludePerfettoModule(
+    Token first_non_space_token) {
+  auto tok = tokenizer_.NextNonWhitespace();
+  auto terminal = tokenizer_.NextTerminal();
+  std::string key = tokenizer_.Substr(tok, terminal).sql();
+
+  if (!ValidateModuleName(key)) {
+    base::StackString<1024> err(
+        "Only alphanumeric characters, dots and underscores allowed in include "
+        "keys: '%s'",
+        key.c_str());
+    return ErrorAtToken(tok, err.c_str());
+  }
+
+  statement_ = Include{key};
+  statement_sql_ = tokenizer_.Substr(first_non_space_token, terminal);
+  return true;
+}
+
+bool PerfettoSqlParser::ParseCreatePerfettoTable(Token first_non_space_token) {
   Token table_name = tokenizer_.NextNonWhitespace();
   if (table_name.token_type != SqliteTokenType::TK_ID) {
     base::StackString<1024> err("Invalid table name %.*s",
@@ -170,12 +227,15 @@ bool PerfettoSqlParser::ParseCreatePerfettoTable() {
   }
 
   Token first = tokenizer_.NextNonWhitespace();
-  statement_ = CreateTable{std::move(name),
-                           tokenizer_.Substr(first, tokenizer_.NextTerminal())};
+  Token terminal = tokenizer_.NextTerminal();
+  statement_ = CreateTable{std::move(name), tokenizer_.Substr(first, terminal)};
+  statement_sql_ = tokenizer_.Substr(first_non_space_token, terminal);
   return true;
 }
 
-bool PerfettoSqlParser::ParseCreatePerfettoFunction(bool replace) {
+bool PerfettoSqlParser::ParseCreatePerfettoFunction(
+    bool replace,
+    Token first_non_space_token) {
   std::string prototype;
   Token function_name = tokenizer_.NextNonWhitespace();
   if (function_name.token_type != SqliteTokenType::TK_ID) {
@@ -234,9 +294,10 @@ bool PerfettoSqlParser::ParseCreatePerfettoFunction(bool replace) {
   }
 
   Token first = tokenizer_.NextNonWhitespace();
-  statement_ = CreateFunction{
-      replace, std::move(prototype), std::move(ret),
-      tokenizer_.Substr(first, tokenizer_.NextTerminal()), table_return};
+  Token terminal = tokenizer_.NextTerminal();
+  statement_ = CreateFunction{replace, std::move(prototype), std::move(ret),
+                              tokenizer_.Substr(first, terminal), table_return};
+  statement_sql_ = tokenizer_.Substr(first_non_space_token, terminal);
   return true;
 }
 

src/trace_processor/perfetto_sql/engine/perfetto_sql_parser.h

@@ -17,6 +17,7 @@
 #ifndef SRC_TRACE_PROCESSOR_PERFETTO_SQL_ENGINE_PERFETTO_SQL_PARSER_H_
 #define SRC_TRACE_PROCESSOR_PERFETTO_SQL_ENGINE_PERFETTO_SQL_PARSER_H_
 
+#include <optional>
 #include <string_view>
 #include <variant>
 
@@ -41,9 +42,7 @@ class PerfettoSqlParser {
  public:
   // Indicates that the specified SQLite SQL was extracted directly from a
   // PerfettoSQL statement and should be directly executed with SQLite.
-  struct SqliteSql {
-    SqlSource sql;
-  };
+  struct SqliteSql {};
   // Indicates that the specified SQL was a CREATE PERFETTO FUNCTION statement
   // with the following parameters.
   struct CreateFunction {
@@ -59,7 +58,13 @@ class PerfettoSqlParser {
     std::string name;
     SqlSource sql;
   };
-  using Statement = std::variant<SqliteSql, CreateFunction, CreateTable>;
+  // Indicates that the specified SQL was a INCLUDE PERFETTO MODULE statement
+  // with the following parameter.
+  struct Include {
+    std::string key;
+  };
+  using Statement =
+      std::variant<SqliteSql, CreateFunction, CreateTable, Include>;
 
   // Creates a new SQL parser with the a block of PerfettoSQL statements.
   // Concretely, the passed string can contain >1 statement.
@@ -80,6 +85,14 @@ class PerfettoSqlParser {
     return statement_.value();
   }
 
+  // Returns the full statement which was parsed. This should return
+  // |statement()| and Perfetto SQL code that's in front. This function *must
+  // not* be called unless |Next()| returned true.
+  const SqlSource& statement_sql() const {
+    PERFETTO_CHECK(statement_sql_);
+    return *statement_sql_;
+  }
+
   // Returns the error status for the parser. This will be |base::OkStatus()|
   // until
   const base::Status& status() const { return status_; }
@@ -89,16 +102,21 @@ class PerfettoSqlParser {
   PerfettoSqlParser(PerfettoSqlParser&&) = delete;
   PerfettoSqlParser& operator=(PerfettoSqlParser&&) = delete;
 
-  bool ParseCreatePerfettoFunction(bool replace);
+  bool ParseCreatePerfettoFunction(
+      bool replace,
+      SqliteTokenizer::Token first_non_space_token);
+
+  bool ParseCreatePerfettoTable(SqliteTokenizer::Token first_non_space_token);
 
-  bool ParseCreatePerfettoTable();
+  bool ParseIncludePerfettoModule(SqliteTokenizer::Token first_non_space_token);
 
   bool ParseArgumentDefinitions(std::string*);
 
   bool ErrorAtToken(const SqliteTokenizer::Token&, const char* error);
 
   SqliteTokenizer tokenizer_;
   base::Status status_;
+  std::optional<SqlSource> statement_sql_;
   std::optional<Statement> statement_;
 };