Merge pull request #363 from mindsdb/create_table

Planner: create empty table
mindsdb · Mar 19, 2024 · ec748c5 · ec748c5
2 parents 727cd00 + cbf499f
commit ec748c5
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 5 deletions.
diff --git a/mindsdb_sql/parser/ast/create.py b/mindsdb_sql/parser/ast/create.py
@@ -16,6 +16,17 @@ def __init__(self, name, type='integer', length=None):
         self.default = None
         self.length = length
 
+    def __eq__(self, other):
+        if type(self) != type(other):
+            return False
+
+        for k in ['name', 'is_primary_key', 'type', 'default', 'length']:
+
+            if getattr(self, k) != getattr(other, k):
+                return False
+
+        return True
+
 
 class CreateTable(ASTNode):
     def __init__(self,

diff --git a/mindsdb_sql/planner/query_planner.py b/mindsdb_sql/planner/query_planner.py
@@ -10,7 +10,7 @@
 from mindsdb_sql.planner.steps import (FetchDataframeStep, ProjectStep, ApplyPredictorStep,
                                        ApplyPredictorRowStep, UnionStep, GetPredictorColumns, SaveToTable,
                                        InsertToTable, UpdateToTable, SubSelectStep,
-                                       DeleteStep, DataStep)
+                                       DeleteStep, DataStep, CreateTableStep)
 from mindsdb_sql.planner.utils import (disambiguate_predictor_column_identifier,
                                        get_deepest_select,
                                        recursively_extract_column_values,
@@ -532,8 +532,16 @@ def plan_project(self, query, dataframe, ignore_doubles=False):
                 out_identifiers.append(new_identifier)
         return self.plan.add_step(ProjectStep(dataframe=dataframe, columns=out_identifiers, ignore_doubles=ignore_doubles))
 
-    def plan_create_table(self, query):
+    def plan_create_table(self, query: CreateTable):
         if query.from_select is None:
+            if query.columns is not None:
+                self.plan.add_step(CreateTableStep(
+                    table=query.name,
+                    columns=query.columns,
+                    is_replace=query.is_replace,
+                ))
+                return
+
             raise PlanningException(f'Not implemented "create table": {query.to_string()}')
 
         integration_name = query.name.parts[0]

diff --git a/mindsdb_sql/planner/steps.py b/mindsdb_sql/planner/steps.py
@@ -227,6 +227,15 @@ def __init__(self, table, dataframe=None, query=None, *args, **kwargs):
         self.query = query
 
 
+class CreateTableStep(PlanStep):
+    def __init__(self, table, columns=None, is_replace=False, *args, **kwargs):
+        """Fills table with content of dataframe"""
+        super().__init__(*args, **kwargs)
+        self.table = table
+        self.columns = columns
+        self.is_replace = is_replace
+
+
 class UpdateToTable(PlanStep):
     def __init__(self, table, dataframe, update_command, *args, **kwargs):
         """Fills table with content of dataframe"""

diff --git a/mindsdb_sql/render/sqlalchemy_render.py b/mindsdb_sql/render/sqlalchemy_render.py
@@ -1,3 +1,4 @@
+import re
 import datetime as dt
 
 import sqlalchemy as sa
@@ -285,8 +286,10 @@ def get_type(self, typename):
             return typename
 
         typename = typename.upper()
-        if typename == 'INT64':
+        if re.match('^INT[\d]*$', typename):
             typename = 'BIGINT'
+        if re.match('^FLOAT[\d]*$', typename):
+            typename = 'FLOAT'
         type = self.types_map[typename]
         return type
 

diff --git a/tests/test_planner/test_integration_select.py b/tests/test_planner/test_integration_select.py
@@ -6,8 +6,7 @@
 from mindsdb_sql.planner import plan_query
 from mindsdb_sql.planner.query_plan import QueryPlan
 from mindsdb_sql.planner.step_result import Result
-from mindsdb_sql.planner.steps import (FetchDataframeStep, ProjectStep, FilterStep, JoinStep, ApplyPredictorStep,
-                                       ApplyPredictorRowStep, GroupByStep, SubSelectStep, UpdateToTable,
+from mindsdb_sql.planner.steps import (FetchDataframeStep, CreateTableStep, SubSelectStep, UpdateToTable,
                                        DeleteStep)
 
 
@@ -689,4 +688,34 @@ def test_delete_from_table_subselect_sql_different_integration(self):
             predictor_metadata=[{'name': 'pred', 'integration_name': 'mindsdb'}]
         )
 
+        assert plan.steps == expected_plan.steps
+
+    def test_create_table(self):
+        query = parse_sql('''
+            CREATE or replace table int2.tab1 (
+            id int8,
+            data varchar
+        ) 
+        ''', dialect='mindsdb')
+
+        expected_plan = QueryPlan(
+            predictor_namespace='mindsdb',
+            steps=[
+                CreateTableStep(
+                    table=Identifier('int2.tab1'),
+                    columns=[
+                        TableColumn(name='id', type='int8'),
+                        TableColumn(name='data', type='varchar'),
+                    ],
+                    is_replace=True
+                ),
+            ],
+        )
+
+        plan = plan_query(
+            query,
+            integrations=[{'name': 'int1', 'class_type': 'api', 'type': 'data'}, 'int2'],
+            predictor_metadata=[{'name': 'pred', 'integration_name': 'mindsdb'}]
+        )
+
         assert plan.steps == expected_plan.steps