Merge pull request #18 from mindsdb/kb-vars

Json can have variable in value

Author: ea-rus

mindsdb_sql_parser/__about__.py

@@ -1,6 +1,6 @@
 __title__ = 'mindsdb_sql_parser'
 __package_name__ = 'mindsdb_sql_parser'
-__version__ = '0.6.1'
+__version__ = '0.7.0'
 __description__ = "Mindsdb SQL parser"
 __email__ = "[email protected]"
 __author__ = 'MindsDB Inc'

mindsdb_sql_parser/ast/variable.py

@@ -15,3 +15,5 @@ def to_tree(self, *args, level=0, **kwargs):
     def get_string(self, *args, **kwargs):
         return ('@@' if self.is_system_var else '@') + f'{str(self.value)}'
 
+    def __repr__(self):
+        return self.get_string()

mindsdb_sql_parser/lexer.py

@@ -350,11 +350,7 @@ def DQUOTE_STRING(self, t):
     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_.$][^"]*"'
-       )
+    @_(r'@(?:([a-zA-Z_$0-9]*[a-zA-Z_$]+[a-zA-Z_$0-9]*)|(?:`([^`]+)`))')
     def VARIABLE(self, t):
         t.value = t.value.lstrip('@')
 
@@ -366,11 +362,7 @@ def VARIABLE(self, t):
             t.value = t.value.strip('`')
         return t
 
-    @_(r'@@[a-zA-Z_.$]+',
-       r"@@'[a-zA-Z_.$][^']*'",
-       r"@@`[a-zA-Z_.$][^`]*`",
-       r'@@"[a-zA-Z_.$][^"]*"'
-       )
+    @_(r'@@(?:([a-zA-Z_$0-9]*[a-zA-Z_$]+[a-zA-Z_$0-9]*)|(?:`([^`]+)`))')
     def SYSTEM_VARIABLE(self, t):
         t.value = t.value.lstrip('@')
 

mindsdb_sql_parser/parser.py

@@ -1716,7 +1716,7 @@ def json_element_list(self, p):
 
     @_('string COLON json_value')
     def json_element(self, p):
-        return {p.string:p.json_value}
+        return {p.string: p.json_value}
 
     # json_array
 
@@ -1736,6 +1736,7 @@ def json_array_list(self, p):
     @_('float',
        'string',
        'integer',
+       'variable',
        'NULL',
        'TRUE',
        'FALSE',

tests/test_mindsdb/test_knowledgebase.py

@@ -1,4 +1,4 @@
-from mindsdb_sql_parser import parse_sql
+from mindsdb_sql_parser import parse_sql, Variable
 from mindsdb_sql_parser.ast.mindsdb.knowledge_base import (
     CreateKnowledgeBase,
     DropKnowledgeBase,
@@ -174,7 +174,7 @@ def test_create_knowledge_base(self):
                 MODEL = mindsdb.my_embedding_model,
                 STORAGE = my_vector_database.some_table,
                 some_param = 'some value',
-                other_param = 'other value'
+                other_param = {'key': @var1}
         """
         ast = parse_sql(sql)
         expected_ast = CreateKnowledgeBase(
@@ -183,7 +183,7 @@ def test_create_knowledge_base(self):
             model=Identifier(parts=["mindsdb", "my_embedding_model"]),
             storage=Identifier(parts=["my_vector_database", "some_table"]),
             from_select=None,
-            params={"some_param": "some value", "other_param": "other value"},
+            params={"some_param": "some value", "other_param": {'key': Variable('var1')}},
         )
         assert ast == expected_ast