Create/drop index for knowledge base (#21)

* adding drop knowledge base index query
* adding create knowledge base index query
* adding tests

Author: fshabashev

mindsdb_sql_parser/ast/mindsdb/knowledge_base.py

@@ -105,3 +105,45 @@ def to_tree(self, *args, level=0, **kwargs):
     def get_string(self, *args, **kwargs):
         out_str = f'DROP KNOWLEDGE_BASE {"IF EXISTS " if self.if_exists else ""}{self.name.to_string()}'
         return out_str
+
+class CreateKnowledgeBaseIndex(ASTNode):
+    """
+    Create a new index in the knowledge base
+    """
+    def __init__(self, name, *args, **kwargs):
+        """
+        Args:
+            name: Identifier -- name of the knowledge base
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+
+    def to_tree(self, *args, level=0, **kwargs):
+        ind = indent(level)
+        out_str = f"{ind}CreateKnowledgeBaseIndex(name={self.name.to_string()})"
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        out_str = f'CREATE INDEX ON KNOWLEDGE_BASE {self.name.to_string()}'
+        return out_str
+
+class DropKnowledgeBaseIndex(ASTNode):
+    """
+    Delete an index in the knowledge base
+    """
+    def __init__(self, name, *args, **kwargs):
+        """
+        Args:
+            name: Identifier -- name of the knowledge base
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+
+    def to_tree(self, *args, level=0, **kwargs):
+        ind = indent(level)
+        out_str = f"{ind}DropKnowledgeBaseIndex(name={self.name.to_string()})"
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        out_str = f'DROP INDEX ON KNOWLEDGE_BASE {self.name.to_string()}'
+        return out_str

mindsdb_sql_parser/parser.py

@@ -16,7 +16,8 @@
 from mindsdb_sql_parser.ast.mindsdb.trigger import CreateTrigger, DropTrigger
 from mindsdb_sql_parser.ast.mindsdb.latest import Latest
 from mindsdb_sql_parser.ast.mindsdb.evaluate import Evaluate
-from mindsdb_sql_parser.ast.mindsdb.knowledge_base import CreateKnowledgeBase, DropKnowledgeBase
+from mindsdb_sql_parser.ast.mindsdb.knowledge_base import CreateKnowledgeBase, DropKnowledgeBase, \
+    CreateKnowledgeBaseIndex, DropKnowledgeBaseIndex
 from mindsdb_sql_parser.ast.mindsdb.skills import CreateSkill, DropSkill, UpdateSkill
 from mindsdb_sql_parser.exceptions import ParsingException
 from mindsdb_sql_parser.ast.mindsdb.retrain_predictor import RetrainPredictor
@@ -93,7 +94,9 @@ class MindsDBParser(Parser):
        'update_skill',
        'create_agent',
        'drop_agent',
-       'update_agent'
+       'update_agent',
+       'create_index',
+       'drop_index',
        )
     def query(self, p):
         return p[0]
@@ -134,6 +137,14 @@ def create_kb(self, p):
             if_not_exists=if_not_exists
         )
 
+    @_('CREATE INDEX ON KNOWLEDGE_BASE identifier')
+    def create_index(self, p):
+        return CreateKnowledgeBaseIndex(name=p.identifier)
+
+    @_('DROP INDEX ON KNOWLEDGE_BASE identifier')
+    def  drop_index(self, p):
+        return DropKnowledgeBaseIndex(name=p.identifier)
+
     @_('DROP KNOWLEDGE_BASE if_exists_or_empty identifier')
     def drop_kb(self, p):
         return DropKnowledgeBase(name=p.identifier, if_exists=p.if_exists_or_empty)

tests/test_mindsdb/test_knowledgebase_index.py

@@ -0,0 +1,29 @@
+from mindsdb_sql_parser import parse_sql
+from mindsdb_sql_parser.ast.mindsdb.knowledge_base import CreateKnowledgeBaseIndex, DropKnowledgeBaseIndex
+from mindsdb_sql_parser.ast import *
+
+class TestKB:
+
+    def test_create_knowledge_base_index(self):
+        # create without select
+
+        sql = """CREATE INDEX ON KNOWLEDGE_BASE my_kb"""
+        ast = parse_sql(sql)
+        print(ast)
+        expected_ast = CreateKnowledgeBaseIndex(
+            name=Identifier('my_kb'),
+        )
+        assert str(ast).lower() == sql.lower()
+        assert ast.to_tree() == expected_ast.to_tree()
+
+    def test_drop_knowledge_base_index(self):
+        # create without select
+
+        sql = """DROP INDEX ON KNOWLEDGE_BASE my_kb"""
+        ast = parse_sql(sql)
+        print(ast)
+        expected_ast = DropKnowledgeBaseIndex(
+            name=Identifier('my_kb'),
+        )
+        assert ast.to_tree() == expected_ast.to_tree()
+