From 8cc2689e38e7ff2a38b3862d118732a710fa2700 Mon Sep 17 00:00:00 2001 From: andrew Date: Mon, 18 Mar 2024 18:19:18 +0300 Subject: [PATCH 1/3] CreateTableStep for empty table --- mindsdb_sql/planner/query_planner.py | 12 ++++++++++-- mindsdb_sql/planner/steps.py | 9 +++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/mindsdb_sql/planner/query_planner.py b/mindsdb_sql/planner/query_planner.py index f8a08b1c..57f37fdd 100644 --- 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 index 7c933f0a..1a1c72ea 100644 --- 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""" From 9be5c83618ffc321babbb0c9e99f5206f4863e93 Mon Sep 17 00:00:00 2001 From: andrew Date: Mon, 18 Mar 2024 18:19:59 +0300 Subject: [PATCH 2/3] render types --- mindsdb_sql/render/sqlalchemy_render.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mindsdb_sql/render/sqlalchemy_render.py b/mindsdb_sql/render/sqlalchemy_render.py index 61637685..22c12dd4 100644 --- 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 From cbf499fe078d3d993e2e3bc7d8d0c940938d541d Mon Sep 17 00:00:00 2001 From: andrew Date: Mon, 18 Mar 2024 18:50:37 +0300 Subject: [PATCH 3/3] create table planner test --- mindsdb_sql/parser/ast/create.py | 11 +++++++ tests/test_planner/test_integration_select.py | 33 +++++++++++++++++-- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/mindsdb_sql/parser/ast/create.py b/mindsdb_sql/parser/ast/create.py index 044eeff6..679b0a0a 100644 --- 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/tests/test_planner/test_integration_select.py b/tests/test_planner/test_integration_select.py index 6cd4e5bf..3cdf26f0 100644 --- 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 \ No newline at end of file