Skip to content

Commit 5649282

Browse files
Merge pull request #885 from Kyrela/2.0
fix: values in compile_alter_sql
2 parents 736eaf8 + ab39abf commit 5649282

12 files changed

+184
-7
lines changed

src/masoniteorm/connections/BaseConnection.py

-1
Original file line numberDiff line numberDiff line change
@@ -85,4 +85,3 @@ def enable_disable_foreign_keys(self):
8585
self._connection.execute(platform.enable_foreign_key_constraints())
8686
elif foreign_keys is not None:
8787
self._connection.execute(platform.disable_foreign_key_constraints())
88-

src/masoniteorm/schema/platforms/MySQLPlatform.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -176,11 +176,16 @@ def compile_alter_sql(self, table):
176176
else:
177177
default = ""
178178

179+
column_constraint = ""
180+
if column.column_type == "enum":
181+
values = ", ".join(f"'{x}'" for x in column.values)
182+
column_constraint = f"({values})"
179183
add_columns.append(
180184
self.add_column_string()
181185
.format(
182186
name=self.get_column_string().format(column=column.name),
183187
data_type=self.type_map.get(column.column_type, ""),
188+
column_constraint=column_constraint,
184189
length=length,
185190
constraint="PRIMARY KEY" if column.primary else "",
186191
nullable="NULL" if column.is_null else "NOT NULL",
@@ -333,14 +338,14 @@ def compile_alter_sql(self, table):
333338

334339
def add_column_string(self):
335340
return (
336-
"ADD {name} {data_type}{length}{signed} {nullable}{default}{after}{comment}"
341+
"ADD {name} {data_type}{length}{column_constraint}{signed} {nullable}{default}{after}{comment}"
337342
)
338343

339344
def drop_column_string(self):
340345
return "DROP COLUMN {name}"
341346

342347
def change_column_string(self):
343-
return "MODIFY {name}{data_type}{length} {nullable}{default} {constraint}"
348+
return "MODIFY {name}{data_type}{length}{column_constraint} {nullable}{default} {constraint}"
344349

345350
def rename_column_string(self):
346351
return "CHANGE {old} {to}"

src/masoniteorm/schema/platforms/PostgresPlatform.py

+15-3
Original file line numberDiff line numberDiff line change
@@ -194,13 +194,19 @@ def compile_alter_sql(self, table):
194194
else:
195195
default = ""
196196

197+
column_constraint = ""
198+
if column.column_type == "enum":
199+
values = ", ".join(f"'{x}'" for x in column.values)
200+
column_constraint = f" CHECK({column.name} IN ({values}))"
201+
197202
add_columns.append(
198203
self.add_column_string()
199204
.format(
200205
name=self.wrap_column(column.name),
201206
data_type=self.type_map.get(column.column_type, ""),
202207
length=length,
203208
constraint="PRIMARY KEY" if column.primary else "",
209+
column_constraint=column_constraint,
204210
nullable="NULL" if column.is_null else "NOT NULL",
205211
default=default,
206212
after=(" AFTER " + self.wrap_column(column._after))
@@ -263,12 +269,18 @@ def compile_alter_sql(self, table):
263269
changed_sql = []
264270

265271
for name, column in table.changed_columns.items():
272+
273+
column_constraint = ""
274+
if column.column_type == "enum":
275+
values = ", ".join(f"'{x}'" for x in column.values)
276+
column_constraint = f" CHECK({column.name} IN ({values}))"
266277
changed_sql.append(
267278
self.modify_column_string()
268279
.format(
269280
name=self.wrap_column(name),
270281
data_type=self.type_map.get(column.column_type),
271-
nullable="NULL" if column.is_null else "NOT NULL",
282+
column_constraint=column_constraint,
283+
constraint="PRIMARY KEY" if column.primary else "",
272284
length="(" + str(column.length) + ")"
273285
if column.column_type not in self.types_without_lengths
274286
else "",
@@ -380,13 +392,13 @@ def alter_format_add_foreign_key(self):
380392
return "ALTER TABLE {table} {columns}"
381393

382394
def add_column_string(self):
383-
return "ADD COLUMN {name} {data_type}{length} {nullable}{default} {constraint}"
395+
return "ADD COLUMN {name} {data_type}{length}{column_constraint} {nullable}{default} {constraint}"
384396

385397
def drop_column_string(self):
386398
return "DROP COLUMN {name}"
387399

388400
def modify_column_string(self):
389-
return "ALTER COLUMN {name} TYPE {data_type}{length}"
401+
return "ALTER COLUMN {name} TYPE {data_type}{length}{column_constraint} {constraint}"
390402

391403
def rename_column_string(self):
392404
return "RENAME COLUMN {old} TO {to}"

src/masoniteorm/schema/platforms/SQLitePlatform.py

+12-1
Original file line numberDiff line numberDiff line change
@@ -171,15 +171,21 @@ def compile_alter_sql(self, diff):
171171
else:
172172
default = ""
173173
constraint = ""
174+
column_constraint = ""
174175
if column.name in diff.added_foreign_keys:
175176
foreign_key = diff.added_foreign_keys[column.name]
176177
constraint = f" REFERENCES {self.wrap_table(foreign_key.foreign_table)}({self.wrap_column(foreign_key.foreign_column)})"
178+
if column.column_type == "enum":
179+
values = ", ".join(f"'{x}'" for x in column.values)
180+
column_constraint = f" CHECK('{column.name}' IN({values}))"
177181

178182
sql.append(
179-
"ALTER TABLE {table} ADD COLUMN {name} {data_type}{signed} {nullable}{default}{constraint}".format(
183+
self.add_column_string()
184+
.format(
180185
table=self.wrap_table(diff.name),
181186
name=self.wrap_column(column.name),
182187
data_type=self.type_map.get(column.column_type, ""),
188+
column_constraint=column_constraint,
183189
nullable="NULL" if column.is_null else "NOT NULL",
184190
default=default,
185191
signed=" " + self.signed.get(column._signed)
@@ -291,6 +297,11 @@ def get_table_string(self):
291297
def get_column_string(self):
292298
return '"{column}"'
293299

300+
def add_column_string(self):
301+
return (
302+
"ALTER TABLE {table} ADD COLUMN {name} {data_type}{column_constraint}{signed} {nullable}{default}{constraint}"
303+
)
304+
294305
def create_column_length(self, column_type):
295306
if column_type in self.types_without_lengths:
296307
return ""

tests/mssql/schema/test_mssql_schema_builder.py

+24
Original file line numberDiff line numberDiff line change
@@ -300,3 +300,27 @@ def test_can_truncate_without_foreign_keys(self):
300300
"ALTER TABLE [users] WITH CHECK CHECK CONSTRAINT ALL",
301301
],
302302
)
303+
304+
def test_can_add_enum(self):
305+
with self.schema.create("users") as blueprint:
306+
blueprint.enum("status", ["active", "inactive"]).default("active")
307+
308+
self.assertEqual(len(blueprint.table.added_columns), 1)
309+
self.assertEqual(
310+
blueprint.to_sql(),
311+
[
312+
"CREATE TABLE [users] ([status] VARCHAR(255) NOT NULL DEFAULT 'active' CHECK([status] IN ('active', 'inactive')))"
313+
],
314+
)
315+
316+
def test_can_change_column_enum(self):
317+
with self.schema.table("users") as blueprint:
318+
blueprint.enum("status", ["active", "inactive"]).default("active").change()
319+
320+
self.assertEqual(len(blueprint.table.changed_columns), 1)
321+
self.assertEqual(
322+
blueprint.to_sql(),
323+
[
324+
"ALTER TABLE [users] ALTER COLUMN [status] VARCHAR(255) NOT NULL DEFAULT 'active' CHECK([status] IN ('active', 'inactive'))"
325+
],
326+
)

tests/mssql/schema/test_mssql_schema_builder_alter.py

+12
Original file line numberDiff line numberDiff line change
@@ -268,3 +268,15 @@ def test_timestamp_alter_add_nullable_column(self):
268268
sql = ["ALTER TABLE [users] ADD [due_date] DATETIME NULL"]
269269

270270
self.assertEqual(blueprint.to_sql(), sql)
271+
272+
def test_can_add_column_enum(self):
273+
with self.schema.table("users") as blueprint:
274+
blueprint.enum("status", ["active", "inactive"]).default("active")
275+
276+
self.assertEqual(len(blueprint.table.added_columns), 1)
277+
278+
sql = [
279+
"ALTER TABLE [users] ADD [status] VARCHAR(255) NOT NULL DEFAULT 'active' CHECK([status] IN ('active', 'inactive'))"
280+
]
281+
282+
self.assertEqual(blueprint.to_sql(), sql)

tests/mysql/schema/test_mysql_schema_builder.py

+12
Original file line numberDiff line numberDiff line change
@@ -384,3 +384,15 @@ def test_can_truncate_without_foreign_keys(self):
384384
"SET FOREIGN_KEY_CHECKS=1",
385385
],
386386
)
387+
388+
def test_can_add_enum(self):
389+
with self.schema.create("users") as blueprint:
390+
blueprint.enum("status", ["active", "inactive"]).default("active")
391+
392+
self.assertEqual(len(blueprint.table.added_columns), 1)
393+
self.assertEqual(
394+
blueprint.to_sql(),
395+
[
396+
"CREATE TABLE `users` (`status` ENUM('active', 'inactive') NOT NULL DEFAULT 'active')"
397+
],
398+
)

tests/mysql/schema/test_mysql_schema_builder_alter.py

+24
Original file line numberDiff line numberDiff line change
@@ -294,3 +294,27 @@ def test_can_create_indexes(self):
294294
"ALTER TABLE `users` ADD FULLTEXT description_fulltext(description)",
295295
],
296296
)
297+
298+
def test_can_add_column_enum(self):
299+
with self.schema.table("users") as blueprint:
300+
blueprint.enum("status", ["active", "inactive"]).default("active")
301+
302+
self.assertEqual(len(blueprint.table.added_columns), 1)
303+
304+
sql = [
305+
"ALTER TABLE `users` ADD `status` ENUM('active', 'inactive') NOT NULL DEFAULT 'active'"
306+
]
307+
308+
self.assertEqual(blueprint.to_sql(), sql)
309+
310+
def test_can_change_column_enum(self):
311+
with self.schema.table("users") as blueprint:
312+
blueprint.enum("status", ["active", "inactive"]).default("active").change()
313+
314+
self.assertEqual(len(blueprint.table.changed_columns), 1)
315+
316+
sql = [
317+
"ALTER TABLE `users` MODIFY `status` ENUM('active', 'inactive') NOT NULL DEFAULT 'active'"
318+
]
319+
320+
self.assertEqual(blueprint.to_sql(), sql)

tests/postgres/schema/test_postgres_schema_builder.py

+12
Original file line numberDiff line numberDiff line change
@@ -368,3 +368,15 @@ def test_can_truncate_without_foreign_keys(self):
368368
'ALTER TABLE "users" ENABLE TRIGGER ALL',
369369
],
370370
)
371+
372+
def test_can_add_enum(self):
373+
with self.schema.create("users") as blueprint:
374+
blueprint.enum("status", ["active", "inactive"]).default("active")
375+
376+
self.assertEqual(len(blueprint.table.added_columns), 1)
377+
self.assertEqual(
378+
blueprint.to_sql(),
379+
[
380+
'CREATE TABLE "users" ("status" VARCHAR(255) CHECK(status IN (\'active\', \'inactive\')) NOT NULL ' 'DEFAULT \'active\')'
381+
],
382+
)

tests/postgres/schema/test_postgres_schema_builder_alter.py

+24
Original file line numberDiff line numberDiff line change
@@ -301,3 +301,27 @@ def test_alter_drop_on_table_schema_table(self):
301301

302302
with schema.table("table_schema") as blueprint:
303303
blueprint.string("name")
304+
305+
def test_can_add_column_enum(self):
306+
with self.schema.table("users") as blueprint:
307+
blueprint.enum("status", ["active", "inactive"]).default("active")
308+
309+
self.assertEqual(len(blueprint.table.added_columns), 1)
310+
311+
sql = [
312+
'ALTER TABLE "users" ADD COLUMN "status" VARCHAR(255) CHECK(status IN (\'active\', \'inactive\')) NOT NULL DEFAULT \'active\'',
313+
]
314+
315+
self.assertEqual(blueprint.to_sql(), sql)
316+
317+
def test_can_change_column_enum(self):
318+
with self.schema.table("users") as blueprint:
319+
blueprint.enum("status", ["active", "inactive"]).default("active").change()
320+
321+
self.assertEqual(len(blueprint.table.changed_columns), 1)
322+
323+
sql = [
324+
'ALTER TABLE "users" ALTER COLUMN "status" TYPE VARCHAR(255) CHECK(status IN (\'active\', \'inactive\')), ALTER COLUMN "status" SET NOT NULL, ALTER COLUMN "status" SET DEFAULT active',
325+
]
326+
327+
self.assertEqual(blueprint.to_sql(), sql)

tests/sqlite/schema/test_sqlite_schema_builder.py

+12
Original file line numberDiff line numberDiff line change
@@ -353,3 +353,15 @@ def test_can_truncate_without_foreign_keys(self):
353353
"PRAGMA foreign_keys = ON",
354354
],
355355
)
356+
357+
def test_can_add_enum(self):
358+
with self.schema.create("users") as blueprint:
359+
blueprint.enum("status", ["active", "inactive"]).default("active")
360+
361+
self.assertEqual(len(blueprint.table.added_columns), 1)
362+
self.assertEqual(
363+
blueprint.to_sql(),
364+
[
365+
'CREATE TABLE "users" ("status" VARCHAR(255) CHECK(status IN (\'active\', \'inactive\')) NOT NULL DEFAULT \'active\')'
366+
],
367+
)

tests/sqlite/schema/test_sqlite_schema_builder_alter.py

+30
Original file line numberDiff line numberDiff line change
@@ -209,3 +209,33 @@ def test_alter_add_foreign_key_only(self):
209209
]
210210

211211
self.assertEqual(blueprint.to_sql(), sql)
212+
213+
def test_can_add_column_enum(self):
214+
with self.schema.table("users") as blueprint:
215+
blueprint.enum("status", ["active", "inactive"]).default("active")
216+
217+
self.assertEqual(len(blueprint.table.added_columns), 1)
218+
219+
sql = [
220+
'ALTER TABLE "users" ADD COLUMN "status" VARCHAR CHECK(\'status\' IN(\'active\', \'inactive\')) NOT NULL DEFAULT \'active\''
221+
]
222+
223+
self.assertEqual(blueprint.to_sql(), sql)
224+
225+
def test_can_change_column_enum(self):
226+
with self.schema.table("users") as blueprint:
227+
blueprint.enum("status", ["active", "inactive"]).default("active").change()
228+
229+
blueprint.table.from_table = Table("users")
230+
231+
self.assertEqual(len(blueprint.table.changed_columns), 1)
232+
233+
sql = [
234+
'CREATE TEMPORARY TABLE __temp__users AS SELECT FROM users',
235+
'DROP TABLE "users"',
236+
'CREATE TABLE "users" ("status" VARCHAR(255) CHECK(status IN (\'active\', \'inactive\')) NOT NULL DEFAULT \'active\')',
237+
'INSERT INTO "users" ("status") SELECT status FROM __temp__users',
238+
'DROP TABLE __temp__users'
239+
]
240+
241+
self.assertEqual(blueprint.to_sql(), sql)

0 commit comments

Comments
 (0)