Add sql syntax for knowledge base

mindsdb_sql/parser/dialects/mindsdb/__init__.py

@@ -16,6 +16,8 @@
 from .drop_job import DropJob
 from .chatbot import CreateChatBot, UpdateChatBot, DropChatBot
 from .trigger import CreateTrigger, DropTrigger
+from .knowledge_base import CreateKnowledgeBase, DropKnowledgeBase
 
 # remove it in next release
 CreateDatasource = CreateDatabase
+

mindsdb_sql/parser/dialects/mindsdb/knowledge_base.py

@@ -0,0 +1,78 @@
+from mindsdb_sql.parser.ast.base import ASTNode
+from mindsdb_sql.parser.utils import indent
+
+
+class CreateKnowledgeBase(ASTNode):
+    def __init__(
+        self,
+        name,
+        model,
+        storage,
+        from_query=None,
+        params=None,
+        if_not_exists=False,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.model = model
+        self.storage = storage
+        self.params = params
+        self.if_not_exists = if_not_exists
+        self.from_query = from_query
+
+    def to_tree(self, *args, level=0, **kwargs):
+        ind = indent(level)
+        out_str = f"""
+        {ind}CreateKnowledgeBase(
+        {ind}    if_not_exists={self.if_not_exists},
+        {ind}    name={self.name.to_string()},
+        {ind}    from_query={self.from_query.to_tree(level=level+1) if self.from_query else None},
+        {ind}    model={self.model.to_string()},
+        {ind}    storage={self.storage.to_string()},
+        {ind}    params={self.params}
+        {ind})
+        """
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        params = self.params.copy()
+        using_ar = [f"{k}={repr(v)}" for k, v in params.items()]
+        using_str = ", ".join(using_ar)
+        from_query_str = (
+            f"FROM ({self.from_query.get_string()})" if self.from_query else ""
+        )
+
+        out_str = (
+            f"CREATE KNOWLEDGE_BASE {'IF NOT EXISTS' if self.if_not_exists else ''}{self.name.to_string()} "
+            f"{from_query_str} "
+            f"MODEL {self.model.to_string()} "
+            f"STORAGE {self.storage.to_string()} "
+            f"USING {using_str}"
+        )
+
+        return out_str
+
+    def __repr__(self) -> str:
+        return self.to_tree()
+
+
+class DropKnowledgeBase(ASTNode):
+    def __init__(self, name, if_exists=False, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.if_exists = if_exists
+
+    def to_tree(self, *args, level=0, **kwargs):
+        ind = indent(level)
+        out_str = (
+            f"{ind}DropKnowledgeBase("
+            f"{ind}    if_exists={self.if_exists},"
+            f"name={self.name.to_string()})"
+        )
+        return out_str
+
+    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

mindsdb_sql/parser/dialects/mindsdb/lexer.py

@@ -29,6 +29,7 @@ class MindsDBLexer(Lexer):
         FINETUNE, EVALUATE,
         LATEST, HORIZON, USING,
         ENGINE, TRAIN, PREDICT, PARAMETERS, JOB, CHATBOT, EVERY,PROJECT,
+        KNOWLEDGE_BASE, KNOWLEDGE_BASES,
 
         # SHOW/DDL Keywords
 
@@ -110,6 +111,9 @@ class MindsDBLexer(Lexer):
     PROJECT = r'\bPROJECT\b'
     EVALUATE = r'\bEVALUATE\b'
 
+    KNOWLEDGE_BASE = r'\bKNOWLEDGE[_|\s]BASE\b'
+    KNOWLEDGE_BASES = r'\bKNOWLEDGE[_|\s]BASES\b'
+
     # Misc
     SET = r'\bSET\b'
     START = r'\bSTART\b'

mindsdb_sql/parser/dialects/mindsdb/parser.py

@@ -16,6 +16,7 @@
 from mindsdb_sql.parser.dialects.mindsdb.latest import Latest
 from mindsdb_sql.parser.dialects.mindsdb.evaluate import Evaluate
 from mindsdb_sql.parser.dialects.mindsdb.create_file import CreateFile
+from mindsdb_sql.parser.dialects.mindsdb.knowledge_base import CreateKnowledgeBase, DropKnowledgeBase
 from mindsdb_sql.exceptions import ParsingException
 from mindsdb_sql.parser.dialects.mindsdb.lexer import MindsDBLexer
 from mindsdb_sql.parser.dialects.mindsdb.retrain_predictor import RetrainPredictor
@@ -79,10 +80,47 @@ class MindsDBParser(Parser):
        'update_chat_bot',
        'create_trigger',
        'drop_trigger',
+       'create_kb',
+       'drop_kb',
        )
     def query(self, p):
         return p[0]
 
+    # -- Knowledge Base --
+    @_(
+        'CREATE KNOWLEDGE_BASE identifier MODEL identifier STORAGE identifier',
+        'CREATE KNOWLEDGE_BASE identifier MODEL identifier STORAGE identifier USING kw_parameter_list',
+        # from select
+        'CREATE KNOWLEDGE_BASE identifier FROM LPAREN select RPAREN MODEL identifier STORAGE identifier',
+        'CREATE KNOWLEDGE_BASE identifier FROM LPAREN select RPAREN MODEL identifier STORAGE identifier USING kw_parameter_list',
+        'CREATE KNOWLEDGE_BASE IF_NOT_EXISTS identifier MODEL identifier STORAGE identifier',
+        'CREATE KNOWLEDGE_BASE IF_NOT_EXISTS identifier MODEL identifier STORAGE identifier USING kw_parameter_list',
+        'CREATE KNOWLEDGE_BASE IF_NOT_EXISTS identifier FROM LPAREN select RPAREN MODEL identifier STORAGE identifier',
+        'CREATE KNOWLEDGE_BASE IF_NOT_EXISTS identifier FROM LPAREN select RPAREN MODEL identifier STORAGE identifier USING kw_parameter_list',
+    )
+    def create_kb(self, p):
+        params = getattr(p, 'kw_parameter_list', {})
+        from_query = getattr(p, 'select', None)
+        name = p.identifier0
+        model = p.identifier1
+        storage = p.identifier2
+        if_not_exists = hasattr(p, 'IF_NOT_EXISTS')
+
+        return CreateKnowledgeBase(
+            name=name,
+            model=model,
+            storage=storage,
+            from_query=from_query,
+            params=params,
+            if_not_exists=if_not_exists
+        )
+
+    @_('DROP KNOWLEDGE_BASE identifier',
+       'DROP KNOWLEDGE_BASE IF_EXISTS identifier')
+    def drop_kb(self, p):
+        if_exists = hasattr(p, 'IF_EXISTS')
+        return DropKnowledgeBase(name=p.identifier, if_exists=if_exists)
+
     # -- ChatBot --
     @_('CREATE CHATBOT identifier USING kw_parameter_list')
     def create_chat_bot(self, p):
@@ -460,6 +498,7 @@ def show(self, p):
        'ML_ENGINES',
        'HANDLERS',
        'SEARCH_PATH',
+       'KNOWLEDGE_BASES',
        'ALL')
     def show_category(self, p):
         return ' '.join([x for x in p])