Fix qlobject->CreateError() by ensuring ast is available in more palces

Summary:
CreateError() wasn't reliable for a long time because sometimes it outputs a proper AST error and sometimes it just returns an invalid argument error.
This diff attempts to fix this by adding a new call inside of LookupVariable to check if ast
is there, adds a dcheck on CreateError for ast ptr.

Using those 2 things, I looked for the objects that were missing asts. I fixed up the probes/mutations objects, collections, and dictionaries to now support ast.

Test Plan: Adjusted tests that assumed the line,col values didn't exist yet.

Reviewers: jamesbartlett, nserrino

Reviewed By: nserrino

Signed-off-by: Phillip Kuznetsov <pkuznetsov@pixielabs.ai>

Differential Revision: https://phab.corp.pixielabs.ai/D11018

GitOrigin-RevId: 53a665d

src/carnot/planner/compiler/ast_visitor.cc

@@ -840,6 +840,7 @@ StatusOr<QLObjectPtr> ASTVisitorImpl::LookupVariable(const pypa::AstPtr& ast,
   if (var == nullptr) {
     return CreateAstError(ast, "name '$0' is not defined", name);
   }
+  var->SetAst(ast);
   return var;
 }
 
@@ -901,14 +902,14 @@ StatusOr<QLObjectPtr> ASTVisitorImpl::ProcessList(const pypa::AstListPtr& ast,
                                                   const OperatorContext& op_context) {
   PL_ASSIGN_OR_RETURN(std::vector<QLObjectPtr> expr_vec,
                       ProcessCollectionChildren(ast->elements, op_context));
-  return ListObject::Create(expr_vec, this);
+  return ListObject::Create(ast, expr_vec, this);
 }
 
 StatusOr<QLObjectPtr> ASTVisitorImpl::ProcessTuple(const pypa::AstTuplePtr& ast,
                                                    const OperatorContext& op_context) {
   PL_ASSIGN_OR_RETURN(std::vector<QLObjectPtr> expr_vec,
                       ProcessCollectionChildren(ast->elements, op_context));
-  return TupleObject::Create(expr_vec, this);
+  return TupleObject::Create(ast, expr_vec, this);
 }
 
 StatusOr<QLObjectPtr> ASTVisitorImpl::ProcessNumber(const pypa::AstNumberPtr& node) {
@@ -1094,7 +1095,7 @@ StatusOr<QLObjectPtr> ASTVisitorImpl::ProcessDict(const pypa::AstDictPtr& node,
     PL_ASSIGN_OR_RETURN(auto value_obj, Process(value, op_context));
     values.push_back(value_obj);
   }
-  return DictObject::Create(keys, values, this);
+  return DictObject::Create(node, keys, values, this);
 }
 
 StatusOr<QLObjectPtr> ASTVisitorImpl::ProcessDataUnaryOp(const pypa::AstUnaryOpPtr& node,

src/carnot/planner/compiler/ast_visitor_test.cc

@@ -1362,9 +1362,8 @@ def plot_latency():
 )pxl";
 
 TEST_F(ASTVisitorTest, problem_decorator_parsed) {
-  auto graph_or_s = CompileGraph(kProblemDecoratorParsing);
-  ASSERT_NOT_OK(graph_or_s);
-  EXPECT_THAT(graph_or_s.status(), HasCompilerError("'vis' object is not callable"));
+  EXPECT_THAT(CompileGraph(kProblemDecoratorParsing).status(),
+              HasCompilerError("'vis' object is not callable"));
 }
 
 constexpr char kGlobalDocStringQuery[] = R"pxl(
@@ -1975,8 +1974,16 @@ df = df.head(n=max_num_records)
 px.display(df)
 )pxl";
 TEST_F(ASTVisitorTest, agg_segfault) {
-  EXPECT_THAT(CompileGraph(kAggSegfaultScript).status().msg(),
-              ::testing::ContainsRegex(".*?All elements of the agg tuple must be column names.*?"));
+  EXPECT_THAT(CompileGraph(kAggSegfaultScript).status(),
+              HasCompilerError(".*?All elements of the agg tuple must be column names.*?"));
+}
+
+TEST_F(ASTVisitorTest, error_on_global) {
+  // Tests to make sure that we will have a valid error on global objects that don't start with ast
+  // errors. AstVisitorImpl::LookupVariable should set the ast for any referenced object.
+  auto vt = VarTable::Create();
+  EXPECT_OK(ParseScript(vt, "l = list"));
+  EXPECT_THAT(vt->Lookup("l")->CreateError("new error"), HasCompilerErrorAt(1, 5, "new error"));
 }
 
 }  // namespace compiler

src/carnot/planner/compiler/test_utils.h

@@ -1129,6 +1129,35 @@ class ASTVisitorTest : public OperatorTests {
     return CompileGraph(query, /* exec_funcs */ {}, /* module_name_to_pxl */ {});
   }
 
+  /**
+   * @brief ParseScript takes a script and an initial variable state then parses
+   * and walks through the AST given this initial variable state, updating the state
+   * with whatever was walked through.
+   *
+   * If you're testing Objects, create a var_table, fill it with the object(s) you want
+   * to test, then write a script interacting with those objects and store the result
+   * into a variable.
+   *
+   * Then you can check if the ParseScript actually succeeds and what data is stored
+   * in the var table.
+   */
+  Status ParseScript(const std::shared_ptr<compiler::VarTable>& var_table,
+                     const std::string& script) {
+    Parser parser;
+    PL_ASSIGN_OR_RETURN(pypa::AstModulePtr ast, parser.Parse(script));
+
+    bool func_based_exec = false;
+    absl::flat_hash_set<std::string> reserved_names;
+    compiler::ModuleHandler module_handler;
+    compiler::MutationsIR mutations_ir;
+    PL_ASSIGN_OR_RETURN(auto ast_walker,
+                        compiler::ASTVisitorImpl::Create(graph.get(), var_table, &mutations_ir,
+                                                         compiler_state_.get(), &module_handler,
+                                                         func_based_exec, reserved_names));
+
+    return ast_walker->ProcessModuleNode(ast);
+  }
+
   StatusOr<std::shared_ptr<IR>> CompileGraph(
       const std::string& query, const std::vector<plannerpb::FuncToExecute>& exec_funcs,
       const absl::flat_hash_map<std::string, std::string>& module_name_to_pxl = {}) {

src/carnot/planner/objects/collection_object.h

@@ -33,14 +33,14 @@ namespace compiler {
 class CollectionObject : public QLObject {
  public:
   const std::vector<QLObjectPtr>& items() const { return *items_; }
-  static bool IsCollection(QLObjectPtr obj) {
+  static bool IsCollection(const QLObjectPtr& obj) {
     return obj->type() == QLObjectType::kList || obj->type() == QLObjectType::kTuple;
   }
 
  protected:
-  CollectionObject(const std::vector<QLObjectPtr>& items, const TypeDescriptor& td,
-                   ASTVisitor* visitor)
-      : QLObject(td, visitor) {
+  CollectionObject(const pypa::AstPtr& ast, const std::vector<QLObjectPtr>& items,
+                   const TypeDescriptor& td, ASTVisitor* visitor)
+      : QLObject(td, ast, visitor) {
     items_ = std::make_shared<std::vector<QLObjectPtr>>(items);
   }
 
@@ -67,16 +67,17 @@ class TupleObject : public CollectionObject {
       /* type */ QLObjectType::kTuple,
   };
 
-  static StatusOr<std::shared_ptr<TupleObject>> Create(const std::vector<QLObjectPtr>& items,
+  static StatusOr<std::shared_ptr<TupleObject>> Create(const pypa::AstPtr& ast,
+                                                       const std::vector<QLObjectPtr>& items,
                                                        ASTVisitor* visitor) {
-    auto tuple = std::shared_ptr<TupleObject>(new TupleObject(items, visitor));
+    auto tuple = std::shared_ptr<TupleObject>(new TupleObject(ast, items, visitor));
     PL_RETURN_IF_ERROR(tuple->Init());
     return tuple;
   }
 
  protected:
-  TupleObject(const std::vector<QLObjectPtr>& items, ASTVisitor* visitor)
-      : CollectionObject(items, TupleType, visitor) {}
+  TupleObject(const pypa::AstPtr& ast, const std::vector<QLObjectPtr>& items, ASTVisitor* visitor)
+      : CollectionObject(ast, items, TupleType, visitor) {}
 };
 
 /**
@@ -89,16 +90,17 @@ class ListObject : public CollectionObject {
       /* type */ QLObjectType::kList,
   };
 
-  static StatusOr<std::shared_ptr<ListObject>> Create(const std::vector<QLObjectPtr>& items,
+  static StatusOr<std::shared_ptr<ListObject>> Create(const pypa::AstPtr& ast,
+                                                      const std::vector<QLObjectPtr>& items,
                                                       ASTVisitor* visitor) {
-    auto list = std::shared_ptr<ListObject>(new ListObject(items, visitor));
+    auto list = std::shared_ptr<ListObject>(new ListObject(ast, items, visitor));
     PL_RETURN_IF_ERROR(list->Init());
     return list;
   }
 
  protected:
-  ListObject(const std::vector<QLObjectPtr>& items, ASTVisitor* visitor)
-      : CollectionObject(items, ListType, visitor) {}
+  ListObject(const pypa::AstPtr& ast, const std::vector<QLObjectPtr>& items, ASTVisitor* visitor)
+      : CollectionObject(ast, items, ListType, visitor) {}
 };
 
 /**

src/carnot/planner/objects/dict_object.h

@@ -39,19 +39,20 @@ class DictObject : public QLObject {
 
   static bool IsDict(const QLObjectPtr& object) { return object->type() == DictType.type(); }
 
-  static StatusOr<std::shared_ptr<DictObject>> Create(const std::vector<QLObjectPtr>& keys,
+  static StatusOr<std::shared_ptr<DictObject>> Create(const pypa::AstPtr& ast,
+                                                      const std::vector<QLObjectPtr>& keys,
                                                       const std::vector<QLObjectPtr>& values,
                                                       ASTVisitor* visitor) {
-    return std::shared_ptr<DictObject>(new DictObject(keys, values, visitor));
+    return std::shared_ptr<DictObject>(new DictObject(ast, keys, values, visitor));
   }
 
   const std::vector<QLObjectPtr>& keys() const { return *keys_; }
   const std::vector<QLObjectPtr>& values() const { return *values_; }
 
  protected:
-  DictObject(const std::vector<QLObjectPtr>& keys, const std::vector<QLObjectPtr>& values,
-             ASTVisitor* visitor)
-      : QLObject(DictType, visitor) {
+  DictObject(const pypa::AstPtr& ast, const std::vector<QLObjectPtr>& keys,
+             const std::vector<QLObjectPtr>& values, ASTVisitor* visitor)
+      : QLObject(DictType, ast, visitor) {
     keys_ = std::make_shared<std::vector<QLObjectPtr>>(keys);
     values_ = std::make_shared<std::vector<QLObjectPtr>>(values);
   }

src/carnot/planner/objects/funcobject.cc

@@ -118,6 +118,7 @@ StatusOr<ParsedArgs> FuncObject::PrepareArgs(const ArgMap& args, const pypa::Ast
       continue;
     }
     PL_ASSIGN_OR_RETURN(auto default_node, GetDefault(arg));
+    default_node->SetAst(ast);
     parsed_args.SubDefaultArg(arg, default_node);
   }
 
@@ -144,7 +145,7 @@ StatusOr<QLObjectPtr> FuncObject::GetDefault(std::string_view arg) {
                                                         /*import_px*/ true);
 }
 
-std::string FuncObject::FormatArguments(const absl::flat_hash_set<std::string> args) {
+std::string FuncObject::FormatArguments(const absl::flat_hash_set<std::string>& args) {
   // Joins the argument names by commas and surrounds each arg with single quotes.
   return absl::StrJoin(args, ",", [](std::string* out, const std::string& arg) {
     absl::StrAppend(out, absl::Substitute("'$0'", arg));
@@ -156,21 +157,20 @@ Status FuncObject::AddVisSpec(std::unique_ptr<VisSpec> vis_spec) {
   return Status::OK();
 }
 
-StatusOr<std::shared_ptr<FuncObject>> GetCallMethod(const pypa::AstPtr& ast, QLObjectPtr pyobject) {
+StatusOr<std::shared_ptr<FuncObject>> GetCallMethod(const pypa::AstPtr& ast,
+                                                    const QLObjectPtr& pyobject) {
   std::shared_ptr<FuncObject> func_object;
   if (pyobject->type_descriptor().type() == QLObjectType::kFunction) {
+    pyobject->SetAst(ast);
     return std::static_pointer_cast<FuncObject>(pyobject);
   }
-  auto func_object_or_s = pyobject->GetCallMethod();
-  if (!func_object_or_s.ok()) {
-    return WrapAstError(ast, func_object_or_s.status());
-  }
-  return func_object_or_s.ConsumeValueOrDie();
+  PL_ASSIGN_OR_RETURN(func_object, pyobject->GetCallMethod());
+  func_object->SetAst(ast);
+  return func_object;
 }
 
-
 Status FuncObject::ResolveArgAnnotationsToTypes(
-    const absl::flat_hash_map<std::string, QLObjectPtr> arg_annotation_objs) {
+    const absl::flat_hash_map<std::string, QLObjectPtr>& arg_annotation_objs) {
   for (const auto& [name, obj] : arg_annotation_objs) {
     if (obj->type() != QLObjectType::kType) {
       PL_ASSIGN_OR_RETURN(auto type_obj,

src/carnot/planner/objects/funcobject.h

@@ -184,7 +184,7 @@ class FuncObject : public QLObject {
   VisSpec* vis_spec() const { return vis_spec_.get(); }
 
   Status ResolveArgAnnotationsToTypes(
-      const absl::flat_hash_map<std::string, QLObjectPtr> arg_annotation_objs);
+      const absl::flat_hash_map<std::string, QLObjectPtr>& arg_annotation_objs);
 
   const absl::flat_hash_map<std::string, std::shared_ptr<TypeObject>>& arg_types() const {
     return arg_types_;
@@ -202,7 +202,7 @@ class FuncObject : public QLObject {
 
   bool HasArgType(std::string_view arg);
 
-  std::string FormatArguments(const absl::flat_hash_set<std::string> args);
+  std::string FormatArguments(const absl::flat_hash_set<std::string>& args);
 
   int64_t NumArgs() const { return arguments_.size(); }
   int64_t NumPositionalArgs() const { return NumArgs() - defaults_.size(); }
@@ -224,7 +224,7 @@ class FuncObject : public QLObject {
 };
 
 template <typename TIRNode>
-StatusOr<TIRNode*> GetArgAs(QLObjectPtr arg, std::string_view arg_name) {
+StatusOr<TIRNode*> GetArgAs(const QLObjectPtr& arg, std::string_view arg_name) {
   if (!ExprObject::IsExprObject(arg)) {
     return arg->CreateError("Expected '$0' in arg '$1', got '$2'", IRNodeTraits<TIRNode>::name,
                             arg_name, QLObjectTypeString(arg->type()));
@@ -233,7 +233,8 @@ StatusOr<TIRNode*> GetArgAs(QLObjectPtr arg, std::string_view arg_name) {
 }
 
 template <typename TIRNode>
-StatusOr<TIRNode*> GetArgAs(const pypa::AstPtr& ast, QLObjectPtr arg, std::string_view arg_name) {
+StatusOr<TIRNode*> GetArgAs(const pypa::AstPtr& ast, const QLObjectPtr& arg,
+                            std::string_view arg_name) {
   return WrapError(ast, GetArgAs<TIRNode>(arg, arg_name));
 }
 
@@ -252,7 +253,8 @@ StatusOr<TIRNode*> GetArgAs(const pypa::AstPtr& ast, const ParsedArgs& args,
  * @param pyobject
  * @return StatusOr<std::shared_ptr<FuncObject>>
  */
-StatusOr<std::shared_ptr<FuncObject>> GetCallMethod(const pypa::AstPtr& ast, QLObjectPtr pyobject);
+StatusOr<std::shared_ptr<FuncObject>> GetCallMethod(const pypa::AstPtr& ast,
+                                                    const QLObjectPtr& pyobject);
 
 }  // namespace compiler
 }  // namespace planner

src/carnot/planner/objects/none_object_test.cc

@@ -31,10 +31,9 @@ namespace compiler {
 using ::testing::ElementsAre;
 class NoneObjectTest : public QLObjectTest {};
 TEST_F(NoneObjectTest, TestNoMethodsWork) {
-  std::shared_ptr<NoneObject> none = std::make_shared<NoneObject>(ast_visitor.get());
-  auto status = none->GetMethod("agg");
-  ASSERT_NOT_OK(status);
-  EXPECT_EQ("'None' object has no attribute 'agg'", status.status().msg());
+  std::shared_ptr<NoneObject> none = std::make_shared<NoneObject>(ast, ast_visitor.get());
+  EXPECT_THAT(none->GetMethod("agg").status(),
+              HasCompilerError("'None' object has no attribute 'agg'"));
 }
 
 }  // namespace compiler

src/carnot/planner/objects/ql_object_test.cc

@@ -40,7 +40,8 @@ class TestQLObject : public QLObject {
       /* type */ QLObjectType::kMisc,
   };
 
-  explicit TestQLObject(ASTVisitor* visitor) : QLObject(TestQLObjectType, visitor) {
+  TestQLObject(const pypa::AstPtr& ast, ASTVisitor* visitor)
+      : QLObject(TestQLObjectType, ast, visitor) {
     std::shared_ptr<FuncObject> func_obj =
         FuncObject::Create("func", {}, {}, /* has_variable_len_args */ false,
                            /* has_variable_len_kwargs */ false,
@@ -51,14 +52,15 @@ class TestQLObject : public QLObject {
     AddMethod("func", func_obj);
   }
 
-  StatusOr<QLObjectPtr> SimpleFunc(const pypa::AstPtr&, const ParsedArgs&, ASTVisitor* visitor) {
-    auto out_obj = std::make_shared<TestQLObject>(visitor);
+  StatusOr<QLObjectPtr> SimpleFunc(const pypa::AstPtr& ast, const ParsedArgs&,
+                                   ASTVisitor* visitor) {
+    auto out_obj = std::make_shared<TestQLObject>(ast, visitor);
     return StatusOr<QLObjectPtr>(out_obj);
   }
 };
 
 TEST_F(QLObjectTest, GetMethod) {
-  auto test_object = std::make_shared<TestQLObject>(ast_visitor.get());
+  auto test_object = std::make_shared<TestQLObject>(ast, ast_visitor.get());
   // No attributes in the TestQLObject, so should return the method
   EXPECT_TRUE(test_object->HasMethod("func"));
   auto attr_or_s = test_object->GetMethod("func");
@@ -74,7 +76,7 @@ TEST_F(QLObjectTest, GetMethod) {
 
 // Gets the method via an attribute.
 TEST_F(QLObjectTest, GetMethodAsAttribute) {
-  auto test_object = std::make_shared<TestQLObject>(ast_visitor.get());
+  auto test_object = std::make_shared<TestQLObject>(ast, ast_visitor.get());
   // No attributes in the TestQLObject, so should return the method
   EXPECT_TRUE(test_object->HasAttribute("func"));
   auto attr_or_s = test_object->GetAttribute(ast, "func");
@@ -91,7 +93,7 @@ TEST_F(QLObjectTest, GetMethodAsAttribute) {
 
 // Attribute not found in either impl or the methods.
 TEST_F(QLObjectTest, AttributeNotFound) {
-  auto test_object = std::make_shared<TestQLObject>(ast_visitor.get());
+  auto test_object = std::make_shared<TestQLObject>(ast, ast_visitor.get());
   // No attributes in the TestQLObject, so should return the method
   std::string attr_name = "bar";
   auto attr_or_s = test_object->GetAttribute(ast, attr_name);
@@ -103,34 +105,32 @@ TEST_F(QLObjectTest, AttributeNotFound) {
 
 // Method not found.
 TEST_F(QLObjectTest, MethodNotFound) {
-  auto test_object = std::make_shared<TestQLObject>(ast_visitor.get());
+  auto test_object = std::make_shared<TestQLObject>(ast, ast_visitor.get());
   // No attributes in the TestQLObject, so should return the method
   std::string attr_name = "bar";
   auto attr_or_s = test_object->GetMethod(attr_name);
   ASSERT_NOT_OK(attr_or_s);
-  EXPECT_THAT(attr_or_s.status().msg(),
-              ContainsRegex(absl::Substitute("'$0'.* has no attribute .*$1",
-                                             test_object->type_descriptor().name(), attr_name)));
+  EXPECT_THAT(attr_or_s.status(),
+              HasCompilerError("'$0'.* has no attribute .*$1",
+                               test_object->type_descriptor().name(), attr_name));
 }
 
 TEST_F(QLObjectTest, NotSubscriptable) {
-  auto test_object = std::make_shared<TestQLObject>(ast_visitor.get());
+  auto test_object = std::make_shared<TestQLObject>(ast, ast_visitor.get());
   // No attributes in the TestQLObject, so should return the method
   auto attr_or_s = test_object->GetSubscriptMethod();
   ASSERT_NOT_OK(attr_or_s);
-  EXPECT_THAT(attr_or_s.status().msg(),
-              ContainsRegex(absl::Substitute("'$0'.* is not subscriptable",
-                                             test_object->type_descriptor().name())));
+  EXPECT_THAT(attr_or_s.status(), HasCompilerError("'$0'.* is not subscriptable",
+                                                   test_object->type_descriptor().name()));
 }
 
 TEST_F(QLObjectTest, NotCallable) {
-  auto test_object = std::make_shared<TestQLObject>(ast_visitor.get());
+  auto test_object = std::make_shared<TestQLObject>(ast, ast_visitor.get());
   // No attributes in the TestQLObject, so should return the method
   auto attr_or_s = test_object->GetCallMethod();
   ASSERT_NOT_OK(attr_or_s);
-  EXPECT_THAT(attr_or_s.status().msg(),
-              ContainsRegex(absl::Substitute("'$0'.* is not callable",
-                                             test_object->type_descriptor().name())));
+  EXPECT_THAT(attr_or_s.status(),
+              HasCompilerError("'$0'.* is not callable", test_object->type_descriptor().name()));
 }
 
 class TestQLObject2 : public QLObject {