variable assignment implemented

Author: wumpz

README.md

@@ -68,6 +68,7 @@ Also I would like to know about needed examples or documentation stuff.
 ## Extensions in the latest SNAPSHOT version 3.3
 
 * allowed Jdbc named parameters within interval expressions
+* allow variable assignments in select statements in different positions
 
 ## Extensions of JSqlParser releases
 

src/main/java/net/sf/jsqlparser/expression/ExpressionVisitor.java

@@ -158,4 +158,6 @@ public interface ExpressionVisitor {
 
     public void visit(ArrayExpression aThis);
 
+    public void visit(VariableAssignment aThis);
+
 }

src/main/java/net/sf/jsqlparser/expression/ExpressionVisitorAdapter.java

@@ -541,4 +541,10 @@ public void visit(ArrayExpression array) {
         array.getObjExpression().accept(this);
         array.getIndexExpression().accept(this);
     }
+
+    @Override
+    public void visit(VariableAssignment var) {
+        var.getVariable().accept(this);
+        var.getExpression().accept(this);
+    }
 }

src/main/java/net/sf/jsqlparser/expression/VariableAssignment.java

@@ -0,0 +1,58 @@
+/*-
+ * #%L
+ * JSQLParser library
+ * %%
+ * Copyright (C) 2004 - 2019 JSQLParser
+ * %%
+ * Dual licensed under GNU LGPL 2.1 or Apache License 2.0
+ * #L%
+ */
+package net.sf.jsqlparser.expression;
+
+import net.sf.jsqlparser.parser.ASTNodeAccessImpl;
+
+/**
+ * Assignment to a user variable like in select @a = 5.
+ */
+public class VariableAssignment extends ASTNodeAccessImpl implements Expression {
+
+    private UserVariable variable;
+    private String operation;
+    private Expression expression;
+
+    public UserVariable getVariable() {
+        return variable;
+    }
+
+    public void setVariable(UserVariable variable) {
+        this.variable = variable;
+    }
+
+    public String getOperation() {
+        return operation;
+    }
+
+    public void setOperation(String operation) {
+        this.operation = operation;
+    }
+
+    public Expression getExpression() {
+        return expression;
+    }
+
+    public void setExpression(Expression expression) {
+        this.expression = expression;
+    }
+
+    @Override
+    public String toString() {
+        return variable.toString() + " " + operation + " " + expression.toString();
+    }
+
+    @Override
+    public void accept(ExpressionVisitor expressionVisitor) {
+        expressionVisitor.visit(this);
+    }
+
+    
+}

src/main/java/net/sf/jsqlparser/util/TablesNamesFinder.java

@@ -49,6 +49,7 @@
 import net.sf.jsqlparser.expression.TimestampValue;
 import net.sf.jsqlparser.expression.UserVariable;
 import net.sf.jsqlparser.expression.ValueListExpression;
+import net.sf.jsqlparser.expression.VariableAssignment;
 import net.sf.jsqlparser.expression.WhenClause;
 import net.sf.jsqlparser.expression.operators.arithmetic.*;
 import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
@@ -891,4 +892,10 @@ public void visit(CreateFunctionalStatement createFunctionalStatement) {
     public void visit(ShowTablesStatement showTables) {
         throw new UnsupportedOperationException("Finding tables from ShowTablesStatement is not supported");
     }
+    
+    @Override
+    public void visit(VariableAssignment var) {
+        var.getVariable().accept(this);
+        var.getExpression().accept(this);
+    }
 }

src/main/java/net/sf/jsqlparser/util/deparser/ExpressionDeParser.java

@@ -50,6 +50,7 @@
 import net.sf.jsqlparser.expression.TimestampValue;
 import net.sf.jsqlparser.expression.UserVariable;
 import net.sf.jsqlparser.expression.ValueListExpression;
+import net.sf.jsqlparser.expression.VariableAssignment;
 import net.sf.jsqlparser.expression.WhenClause;
 import net.sf.jsqlparser.expression.WindowElement;
 import net.sf.jsqlparser.expression.operators.arithmetic.Addition;
@@ -893,4 +894,11 @@ public void visit(ArrayExpression array) {
     void deParse(Expression statement) {
         statement.accept(this);
     }
+    
+    @Override
+    public void visit(VariableAssignment var) {
+        var.getVariable().accept(this);
+        buffer.append(" ").append(var.getOperation()).append(" ");
+        var.getExpression().accept(this);
+    }
 }

src/main/jjtree/net/sf/jsqlparser/parser/JSqlParserCC.jjt

@@ -3030,12 +3030,24 @@ Expression AnyComparisonExpression() :
 Expression SimpleExpression():
 {
     Expression retval = null;
+    UserVariable user = null;
+    Token operation = null;
 }
 {
-    [ LOOKAHEAD(UserVariable() "=") UserVariable() "=" ]
+    [ LOOKAHEAD(UserVariable() ("=" | ":=") ) 
+        user = UserVariable() 
+       ( operation = "=" | operation = ":=" )
+    ]
     retval=ConcatExpression()
     {
-        return retval;
+        if (user != null) {
+            VariableAssignment assignment = new VariableAssignment();
+            assignment.setVariable(user);
+            assignment.setOperation(operation.image);
+            assignment.setExpression(retval);
+            return assignment;
+        } else 
+            return retval;
     }
 }
 

src/test/java/net/sf/jsqlparser/statement/select/SelectTest.java

@@ -4185,4 +4185,9 @@ public void testCheckDateFunctionIssue_3() throws JSQLParserException {
     public void testCheckColonVariable() throws JSQLParserException {
         assertSqlCanBeParsedAndDeparsed("SELECT * FROM mytable WHERE (col1, col2) IN ((:qp0, :qp1), (:qp2, :qp3))");
     }
+    
+    @Test
+    public void testVariableAssignment() throws JSQLParserException {
+        assertSqlCanBeParsedAndDeparsed("SELECT @SELECTVariable = 2");
+    }
 }

src/test/java/net/sf/jsqlparser/statement/update/UpdateTest.java

@@ -160,4 +160,9 @@ public void testUpdateIssue750() throws JSQLParserException {
     public void testUpdateIssue962Validate() throws JSQLParserException {
         assertSqlCanBeParsedAndDeparsed("UPDATE tbl_user_card SET validate = '1', identityCodeFlag = 1 WHERE id = 9150000293816");
     }
+    
+    @Test
+    public void testUpdateVariableAssignment() throws JSQLParserException {
+        assertSqlCanBeParsedAndDeparsed("UPDATE transaction_id SET latest_id_wallet = (@cur_id_wallet := latest_id_wallet) + 1");
+    }
 }