add support for PREDICT clause:

mindsdb · Sep 2, 2023 · 17f39d9 · 17f39d9
1 parent 5f86a6c
commit 17f39d9
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 23 deletions.
diff --git a/mindsdb_sql/parser/dialects/mindsdb/parser.py b/mindsdb_sql/parser/dialects/mindsdb/parser.py
@@ -718,10 +718,10 @@ def create_predictor(self, p):
     ## Anomaly detection
     @_(
         'CREATE ANOMALY DETECTION MODEL identifier',  # for methods that do not require training (e.g. TimeGPT)
-        'CREATE ANOMALY DETECTION MODEL IF_NOT_EXISTS identifier',
         'CREATE ANOMALY DETECTION MODEL identifier FROM identifier LPAREN raw_query RPAREN',
-        'CREATE ANOMALY DETECTION MODEL identifier IF_NOT_EXISTS FROM identifier LPAREN raw_query RPAREN',
-
+        'CREATE ANOMALY DETECTION MODEL identifier PREDICT result_columns',
+        'CREATE ANOMALY DETECTION MODEL identifier PREDICT result_columns FROM identifier LPAREN raw_query RPAREN',
+        # TODO add IF_NOT_EXISTS elegantly (should be low level BNF expansion)
     )
     def create_anomaly_detection_model(self, p):
 
@@ -737,6 +737,7 @@ def create_anomaly_detection_model(self, p):
 
         return CreateAnomalyDetectionModel(
             name=name,
+            targets=getattr(p, 'result_columns', None),
             integration_name=getattr(p, 'identifier1', None),
             query_str=query_str,
             if_not_exists=hasattr(p, 'IF_NOT_EXISTS')

diff --git a/tests/test_parser/test_mindsdb/test_create_predictor.py b/tests/test_parser/test_mindsdb/test_create_predictor.py
@@ -162,23 +162,25 @@ def test_create_predictor_empty_fields(self):
         assert ast.to_tree() == expected_ast.to_tree()
 
     def test_create_anomaly_detection_model(self):
-        sql = """
-        CREATE ANOMALY DETECTION MODEL alert_model
-        FROM integration_name (select * FROM table)
-        USING
-            confidence=0.5
-        """
-        ast = parse_sql(sql, dialect='mindsdb')
-
-        expected_ast = CreateAnomalyDetectionModel(
-            name=Identifier('alert_model'),
-            integration_name=Identifier('integration_name'),
-            query_str='select * FROM table',
-            targets=None,
-            using={
-                'confidence': 0.5
-            }
-        )
-
-        assert to_single_line(str(ast)) == to_single_line(str(expected_ast))
-        assert ast.to_tree() == expected_ast.to_tree()
+        for predict_clause in ["", " PREDICT alert "]:
+            create_clause = """CREATE ANOMALY DETECTION MODEL alert_model """
+            rest_clause = """
+            FROM integration_name (select * FROM table)
+            USING
+                confidence=0.5
+            """
+            sql = create_clause + predict_clause + rest_clause
+            ast = parse_sql(sql, dialect='mindsdb')
+
+            expected_ast = CreateAnomalyDetectionModel(
+                name=Identifier('alert_model'),
+                integration_name=Identifier('integration_name'),
+                query_str='select * FROM table',
+                targets=[Identifier('alert')] if predict_clause else None,
+                using={
+                    'confidence': 0.5
+                }
+            )
+
+            assert to_single_line(str(ast)) == to_single_line(str(expected_ast))
+            assert ast.to_tree() == expected_ast.to_tree()