feat(parser): support set statement for (#38)

- set global xx = 'on' - set session xx = 'off'
4paradigm · Mar 2, 2022 · 31e5647 · 31e5647
1 parent 2e548a0
commit 31e5647
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 1 deletion.
diff --git a/zetasql/parser/ast_node_kind.h b/zetasql/parser/ast_node_kind.h
@@ -322,7 +322,8 @@ enum ASTNodeKind {
   AST_STOP_STATEMENT,
   AST_CONFIG_CLAUSE,
   AST_WITH_WEIGHT,
-  kLastASTNodeKind = AST_WITH_WEIGHT
+  AST_SCOPED_VARIABLE_ASSIGNMENT,
+  kLastASTNodeKind = AST_SCOPED_VARIABLE_ASSIGNMENT
 };
 
 }  // namespace zetasql

diff --git a/zetasql/parser/bison_parser.y b/zetasql/parser/bison_parser.y
@@ -1980,6 +1980,18 @@ set_statement:
     {
       $$ = MAKE_NODE(ASTSystemVariableAssignment, @$, {$2, $4});
     }
+    | "SET" "GLOBAL" identifier "=" expression
+    {
+      auto* stmt = MAKE_NODE(ASTScopedVariableAssignment, @$, {$3, $5});
+      stmt->set_scope(zetasql::ASTScopedVariableAssignment::Scope::GLOBAL);
+      $$ = stmt;
+    }
+    | "SET" "SESSION" identifier "=" expression
+    {
+      auto* stmt = MAKE_NODE(ASTScopedVariableAssignment, @$, {$3, $5});
+      stmt->set_scope(zetasql::ASTScopedVariableAssignment::Scope::SESSION);
+      $$ = stmt;
+    }
     | "SET" "(" identifier_list ")" "=" expression
     {
       $$ = MAKE_NODE(ASTAssignmentFromStruct, @$, {$3, $6});
@@ -8895,6 +8907,10 @@ next_statement_kind_without_hint:
       { $$ = zetasql::ASTParameterAssignment::kConcreteNodeKind; }
     | "SET" system_variable_expression "="
       { $$ = zetasql::ASTSystemVariableAssignment::kConcreteNodeKind; }
+    | "SET" "GLOBAL" identifier "="
+      { $$ = zetasql::ASTScopedVariableAssignment::kConcreteNodeKind; }
+    | "SET" "SESSION" identifier "="
+      { $$ = zetasql::ASTScopedVariableAssignment::kConcreteNodeKind; }
     | "SET" "("
       { $$ = zetasql::ASTAssignmentFromStruct::kConcreteNodeKind; }
     | "COMMIT" { $$ = zetasql::ASTCommitStatement::kConcreteNodeKind; }

diff --git a/zetasql/parser/parse_tree.cc b/zetasql/parser/parse_tree.cc
@@ -359,6 +359,7 @@ static absl::flat_hash_map<ASTNodeKind, std::string> CreateNodeNamesMap() {
   map[AST_CONFIG_CLAUSE] = "ConfigClause";
   map[AST_WITH_WEIGHT] = "WithWeight";
   map[AST_WITH_PARTITION_COLUMNS_CLAUSE] = "WithPartitionColumnsClause";
+  map[AST_SCOPED_VARIABLE_ASSIGNMENT] = "ScopedVariableAssignment";
   for (int kind = kFirstASTNodeKind; kind <= kLastASTNodeKind;
        ++kind) {
     ZETASQL_DCHECK(zetasql_base::ContainsKey(map, static_cast<ASTNodeKind>(kind)))

diff --git a/zetasql/parser/parse_tree_manual.h b/zetasql/parser/parse_tree_manual.h
@@ -8167,6 +8167,44 @@ class ASTSystemVariableAssignment final : public ASTStatement {
   const ASTExpression* expression_ = nullptr;
 };
 
+// Yet another variable assignment statement
+// has the same meaning with ASTSystemVariableAssignment but in different syntax
+// e.g. SET GLOBAL xxx = 'on' is equivalent with SET @@global.xxx = 'on'
+//      SET SESSION xxx = 'off' is equivalent with SET @@session.xxx = 'off'
+class ASTScopedVariableAssignment final : public ASTStatement {
+ public:
+  enum Scope {
+    SESSION,
+    GLOBAL,
+  };
+  static constexpr ASTNodeKind kConcreteNodeKind =
+      AST_SCOPED_VARIABLE_ASSIGNMENT;
+
+  ASTScopedVariableAssignment() : ASTStatement(kConcreteNodeKind) {}
+  void Accept(ParseTreeVisitor* visitor, void* data) const override;
+  zetasql_base::StatusOr<VisitResult> Accept(
+      NonRecursiveParseTreeVisitor* visitor) const override;
+
+  const ASTIdentifier* variable() const {
+    return variable_;
+  }
+  const ASTExpression* expression() const { return expression_; }
+
+  Scope scope() const { return scope_; }
+  void set_scope(Scope scope) { scope_ = scope; }
+
+ private:
+  void InitFields() final {
+    FieldLoader fl(this);
+    fl.AddRequired(&variable_);
+    fl.AddRequired(&expression_);
+  }
+
+  const ASTIdentifier* variable_ = nullptr;
+  const ASTExpression* expression_ = nullptr;
+  Scope scope_ = SESSION;
+};
+
 // A statement which assigns multiple variables to fields in a struct,
 // which each variable assigned to one field.
 // Example:

diff --git a/zetasql/parser/testdata/set.test b/zetasql/parser/testdata/set.test
@@ -430,6 +430,29 @@ SET @@VAL = (
 )
 ==
 
+
+# ScopedVariableAssignment
+
+SET {{SESSION|GLOBAL}} xxx = 'on'
+--
+ALTERNATION GROUP: SESSION
+--
+ScopedVariableAssignment [0-22]
+  Identifier(xxx) [12-15]
+  StringLiteral('on') [18-22]
+--
+SET SESSION xxx = 'on'
+--
+ALTERNATION GROUP: GLOBAL
+--
+ScopedVariableAssignment [0-21]
+  Identifier(xxx) [11-14]
+  StringLiteral('on') [17-21]
+--
+SET GLOBAL xxx = 'on'
+==
+
+
 # Error cases
 
 # Does not work for setting a parameter to an invalid expression.

diff --git a/zetasql/parser/unparser.cc b/zetasql/parser/unparser.cc
@@ -3254,6 +3254,24 @@ void Unparser::visitASTSystemVariableAssignment(
   node->expression()->Accept(this, data);
 }
 
+void Unparser::visitASTScopedVariableAssignment(
+    const ASTScopedVariableAssignment *node, void *data) {
+  print("SET");
+  switch (node->scope()) {
+    case ASTScopedVariableAssignment::Scope::SESSION: {
+      print("SESSION");
+      break;
+    }
+    case ASTScopedVariableAssignment::Scope::GLOBAL: {
+      print("GLOBAL");
+      break;
+    }
+  }
+  node->variable()->Accept(this, data);
+  print("=");
+  node->expression()->Accept(this, data);
+}
+
 void Unparser::visitASTAssignmentFromStruct(const ASTAssignmentFromStruct* node,
                                   void* data) {
   print("SET");

diff --git a/zetasql/parser/unparser.h b/zetasql/parser/unparser.h
@@ -598,6 +598,8 @@ class Unparser : public ParseTreeVisitor {
                                    void* data) override;
   void visitASTSystemVariableAssignment(const ASTSystemVariableAssignment* node,
                                         void* data) override;
+  void visitASTScopedVariableAssignment(const ASTScopedVariableAssignment* node,
+                                        void* data) override;
   void visitASTSingleAssignment(const ASTSingleAssignment* node,
                                 void* data) override;