CREATE TABLE with columns #357

mindsdb · Mar 11, 2024 · 551cf56 · 551cf56
1 parent ab5f61b
commit 551cf56
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 4 deletions.
diff --git a/mindsdb_sql/parser/ast/create.py b/mindsdb_sql/parser/ast/create.py
@@ -9,11 +9,12 @@
 
 
 class TableColumn():
-    def __init__(self, name, type='integer'):
+    def __init__(self, name, type='integer', length=None):
         self.name = name
         self.type = type
         self.is_primary_key = False
         self.default = None
+        self.length = length
 
 
 class CreateTable(ASTNode):
@@ -72,14 +73,18 @@ def get_string(self, *args, **kwargs):
         if self.columns is not None:
             columns = []
             for col in self.columns:
-                type = str(col.type)
-                if sa_types is not None:
+
+                if not isinstance(col.type, str) and sa_types is not None:
                     if issubclass(col.type, sa_types.Integer):
                         type = 'int'
                     elif issubclass(col.type, sa_types.Float):
                         type = 'float'
                     elif issubclass(col.type, sa_types.Text):
                         type = 'text'
+                else:
+                    type = str(col.type)
+                if col.length is not None:
+                    type = f'{type}({col.length})'
                 columns.append( f'{col.name} {type}')
 
             columns_str = '({})'.format(', '.join(columns))

diff --git a/mindsdb_sql/parser/dialects/mindsdb/parser.py b/mindsdb_sql/parser/dialects/mindsdb/parser.py
@@ -690,6 +690,31 @@ def drop_table(self, p):
         return DropTables(tables=[p.identifier], if_exists=p.if_exists_or_empty)
 
     # create table
+    @_('id id',
+       'id id LPAREN INTEGER RPAREN')
+    def table_column(self, p):
+        return TableColumn(
+            name=p[0],
+            type=p[1],
+            length=getattr(p, 'INTEGER', None)
+        )
+
+    @_('table_column',
+       'table_column_list COMMA table_column')
+    def table_column_list(self, p):
+        items = getattr(p, 'table_column_list', [])
+        items.append(p.table_column)
+        return items
+
+    @_('CREATE replace_or_empty TABLE if_not_exists_or_empty identifier LPAREN table_column_list RPAREN')
+    def create_table(self, p):
+        return CreateTable(
+            name=p.identifier,
+            columns=p.table_column_list,
+            is_replace=getattr(p, 'replace_or_empty', False),
+            if_not_exists=getattr(p, 'if_not_exists_or_empty', False)
+        )
+
     @_(
        'CREATE replace_or_empty TABLE if_not_exists_or_empty identifier select',
        'CREATE replace_or_empty TABLE if_not_exists_or_empty identifier LPAREN select RPAREN',

diff --git a/tests/test_parser/test_base_sql/test_create.py b/tests/test_parser/test_base_sql/test_create.py
@@ -3,9 +3,10 @@
 from mindsdb_sql import parse_sql
 from mindsdb_sql.parser.ast import *
 
+
 @pytest.mark.parametrize('dialect', ['mysql', 'mindsdb'])
 class TestCreate:
-    def test_create(self, dialect):
+    def test_create_from_select(self, dialect):
         expected_ast = CreateTable(
             name=Identifier('int1.model_name'),
             is_replace=True,
@@ -49,3 +50,40 @@ def test_create(self, dialect):
         assert ast.to_tree() == expected_ast.to_tree()
 
 
+class TestCreateMindsdb:
+
+    def test_create(self):
+
+        for is_replace in [True, False]:
+            for if_not_exists in [True, False]:
+
+                expected_ast = CreateTable(
+                    name=Identifier('mydb.Persons'),
+                    is_replace=is_replace,
+                    if_not_exists=if_not_exists,
+                    columns=[
+                        TableColumn(name='PersonID', type='int'),
+                        TableColumn(name='LastName', type='varchar', length=255),
+                        TableColumn(name='FirstName', type='char', length=10),
+                        TableColumn(name='Info', type='json'),
+                        TableColumn(name='City', type='varchar'),
+                    ]
+                )
+                replace_str = 'OR REPLACE' if is_replace else ''
+                exist_str = 'IF NOT EXISTS' if if_not_exists else ''
+
+                sql = f'''
+                 CREATE {replace_str} TABLE {exist_str} mydb.Persons(
+                    PersonID int,
+                    LastName varchar(255),
+                    FirstName char(10),
+                    Info json,
+                    City varchar
+                 )
+                '''
+                print(sql)
+                ast = parse_sql(sql)
+
+                assert str(ast).lower() == str(expected_ast).lower()
+                assert ast.to_tree() == expected_ast.to_tree()
+