variables added to mindsdb dialect

#196

mindsdb_sql/parser/ast/__init__.py

@@ -14,6 +14,6 @@
 from .delete import *
 from .drop import *
 from .create import *
+from .variable import *
 
-from mindsdb_sql.parser.dialects.mysql.variable import Variable
 from mindsdb_sql.parser.dialects.mindsdb.latest import Latest

mindsdb_sql/parser/ast/set.py

@@ -16,7 +16,6 @@ def __init__(self,
 
     def to_tree(self, *args, level=0, **kwargs):
         ind = indent(level)
-        ind1 = indent(level+1)
         category_str = f'category={self.category}, '
         arg_str = f'arg={self.arg.to_tree()},' if self.arg else ''
         if self.params:

mindsdb_sql/parser/dialects/mindsdb/lexer.py

@@ -44,6 +44,8 @@ class MindsDBLexer(Lexer):
         ENGINES, CHARSET, COLLATION, PLUGINS, CHARACTER,
         PERSIST, PERSIST_ONLY, DEFAULT,
         IF_EXISTS, IF_NOT_EXISTS, COLUMNS, FIELDS, COLLATE, SEARCH_PATH,
+        VARIABLE, SYSTEM_VARIABLE,
+
         # SELECT Keywords
         WITH, SELECT, DISTINCT, FROM, WHERE, AS,
         LIMIT, OFFSET, ASC, DESC, NULLS_FIRST, NULLS_LAST,
@@ -314,3 +316,36 @@ def DQUOTE_STRING(self, t):
     @_(r'\n+')
     def ignore_newline(self, t):
         self.lineno += len(t.value)
+
+    @_(r'@[a-zA-Z_.$]+',
+       r"@'[a-zA-Z_.$][^']*'",
+       r"@`[a-zA-Z_.$][^`]*`",
+       r'@"[a-zA-Z_.$][^"]*"'
+       )
+    def VARIABLE(self, t):
+        t.value = t.value.lstrip('@')
+
+        if t.value[0] == '"':
+            t.value = t.value.strip('\"')
+        elif t.value[0] == "'":
+            t.value = t.value.strip('\'')
+        elif t.value[0] == "`":
+            t.value = t.value.strip('`')
+        return t
+
+    @_(r'@@[a-zA-Z_.$]+',
+       r"@@'[a-zA-Z_.$][^']*'",
+       r"@@`[a-zA-Z_.$][^`]*`",
+       r'@@"[a-zA-Z_.$][^"]*"'
+       )
+    def SYSTEM_VARIABLE(self, t):
+        t.value = t.value.lstrip('@')
+
+        if t.value[0] == '"':
+            t.value = t.value.strip('\"')
+        elif t.value[0] == "'":
+            t.value = t.value.strip('\'')
+        elif t.value[0] == "`":
+            t.value = t.value.strip('`')
+        return t
+

mindsdb_sql/parser/dialects/mindsdb/parser.py

@@ -1692,6 +1692,22 @@ def raw_query(self, p):
     def raw_query(self, p):
         return p[0] + p[1]
 
+    @_('variable')
+    def table_or_subquery(self, p):
+        return p.variable
+
+    @_('variable')
+    def expr(self, p):
+        return p.variable
+
+    @_('SYSTEM_VARIABLE')
+    def variable(self, p):
+        return Variable(value=p.SYSTEM_VARIABLE, is_system_var=True)
+
+    @_('VARIABLE')
+    def variable(self, p):
+        return Variable(value=p.VARIABLE)
+
     @_(
         'IF_NOT_EXISTS',
         'empty'

mindsdb_sql/parser/dialects/mysql/__init__.py

@@ -1,2 +1 @@
-from .variable import Variable
 from .show_index import ShowIndex

mindsdb_sql/parser/dialects/mysql/parser.py

@@ -2,7 +2,6 @@
 from mindsdb_sql.parser.parser import SQLParser
 from mindsdb_sql.parser.ast import *
 from mindsdb_sql.parser.dialects.mysql.lexer import MySQLLexer
-from mindsdb_sql.parser.dialects.mysql.variable import Variable
 from mindsdb_sql.exceptions import ParsingException
 from mindsdb_sql.parser.utils import ensure_select_keyword_order, JoinType
 

tests/test_parser/test_mindsdb/test_variables.py

@@ -0,0 +1,37 @@
+from mindsdb_sql import parse_sql
+from mindsdb_sql.parser.ast import *
+from mindsdb_sql.parser.ast import Variable
+
+class TestMDBParser:
+    def test_select_variable(self):
+        sql = 'SELECT @version'
+        ast = parse_sql(sql)
+        expected_ast = Select(targets=[Variable('version')])
+        assert str(ast).lower() == sql.lower()
+        assert str(ast) == str(expected_ast)
+
+        sql = 'SELECT @@version'
+        ast = parse_sql(sql)
+        expected_ast = Select(targets=[Variable('version', is_system_var=True)])
+        assert str(ast).lower() == sql.lower()
+        assert str(ast) == str(expected_ast)
+
+        sql = "set autocommit = 1, sql_mode = concat(@@sql_mode, ',STRICT_TRANS_TABLES')"
+        ast = parse_sql(sql)
+        expected_ast = Set(
+            arg=Tuple([
+                BinaryOperation(op='=', args=[
+                    Identifier('autocommit'), Constant(1)
+                ]),
+                BinaryOperation(op='=', args=[
+                    Identifier('sql_mode'),
+                    Function(op='concat', args=[
+                        Variable('sql_mode', is_system_var=True),
+                        Constant(',STRICT_TRANS_TABLES')
+                    ])
+                ])
+            ])
+        )
+        assert str(ast).lower() == sql.lower()
+        assert str(ast) == str(expected_ast)
+

tests/test_parser/test_mysql/test_mysql_parser.py

@@ -1,6 +1,6 @@
 from mindsdb_sql import parse_sql
 from mindsdb_sql.parser.ast import Select, Identifier, BinaryOperation, Star
-from mindsdb_sql.parser.dialects.mysql import Variable
+from mindsdb_sql.parser.ast import Variable
 from mindsdb_sql.parser.parser import Show
 
 class TestMySQLParser: