Skip to content

Commit

Permalink
fix migrations for tables in schemas (#884)
Browse files Browse the repository at this point in the history
  • Loading branch information
@dantownsend
dantownsend authored Sep 8, 2023
1 parent 022e22a commit d2b2a5a
Show file tree
Hide file tree
Showing 7 changed files with 160 additions and 31 deletions.
3 changes: 3 additions & 0 deletions piccolo/apps/migrations/auto/diffable_table.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,7 @@ def __sub__(self, value: DiffableTable) -> TableDelta:
column_class_name=i.column.__class__.__name__,
column_class=i.column.__class__,
params=i.column._meta.params,
schema=self.schema,
)
for i in sorted(
{ColumnComparison(column=column) for column in self.columns}
Expand All @@ -145,6 +146,7 @@ def __sub__(self, value: DiffableTable) -> TableDelta:
column_name=i.column._meta.name,
db_column_name=i.column._meta.db_column_name,
tablename=value.tablename,
schema=self.schema,
)
for i in sorted(
{ColumnComparison(column=column) for column in value.columns}
Expand Down Expand Up @@ -184,6 +186,7 @@ def __sub__(self, value: DiffableTable) -> TableDelta:
old_params=old_params,
column_class=column.__class__,
old_column_class=existing_column.__class__,
schema=self.schema,
)
)

Expand Down
42 changes: 36 additions & 6 deletions piccolo/apps/migrations/auto/migration_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ class AddColumnClass:
column: Column
table_class_name: str
tablename: str
schema: t.Optional[str]


@dataclass
Expand Down Expand Up @@ -206,13 +207,15 @@ def rename_table(
old_tablename: str,
new_class_name: str,
new_tablename: str,
schema: t.Optional[str] = None,
):
self.rename_tables.append(
RenameTable(
old_class_name=old_class_name,
old_tablename=old_tablename,
new_class_name=new_class_name,
new_tablename=new_tablename,
schema=schema,
)
)

Expand All @@ -225,6 +228,7 @@ def add_column(
column_class_name: str = "",
column_class: t.Optional[t.Type[Column]] = None,
params: t.Dict[str, t.Any] = None,
schema: t.Optional[str] = None,
):
"""
Add a new column to the table.
Expand Down Expand Up @@ -255,6 +259,7 @@ def add_column(
column=column,
tablename=tablename,
table_class_name=table_class_name,
schema=schema,
)
)

Expand All @@ -264,13 +269,15 @@ def drop_column(
tablename: str,
column_name: str,
db_column_name: t.Optional[str] = None,
schema: t.Optional[str] = None,
):
self.drop_columns.append(
DropColumn(
table_class_name=table_class_name,
column_name=column_name,
db_column_name=db_column_name or column_name,
tablename=tablename,
schema=schema,
)
)

Expand All @@ -282,6 +289,7 @@ def rename_column(
new_column_name: str,
old_db_column_name: t.Optional[str] = None,
new_db_column_name: t.Optional[str] = None,
schema: t.Optional[str] = None,
):
self.rename_columns.append(
RenameColumn(
Expand All @@ -291,6 +299,7 @@ def rename_column(
new_column_name=new_column_name,
old_db_column_name=old_db_column_name or old_column_name,
new_db_column_name=new_db_column_name or new_column_name,
schema=schema,
)
)

Expand All @@ -304,6 +313,7 @@ def alter_column(
old_params: t.Dict[str, t.Any] = None,
column_class: t.Optional[t.Type[Column]] = None,
old_column_class: t.Optional[t.Type[Column]] = None,
schema: t.Optional[str] = None,
):
"""
All possible alterations aren't currently supported.
Expand All @@ -322,6 +332,7 @@ def alter_column(
old_params=old_params,
column_class=column_class,
old_column_class=old_column_class,
schema=schema,
)
)

Expand Down Expand Up @@ -403,7 +414,10 @@ async def _run_alter_columns(self, backwards=False):

_Table: t.Type[Table] = create_table_class(
class_name=table_class_name,
class_kwargs={"tablename": alter_columns[0].tablename},
class_kwargs={
"tablename": alter_columns[0].tablename,
"schema": alter_columns[0].schema,
},
)

for alter_column in alter_columns:
Expand Down Expand Up @@ -635,7 +649,10 @@ async def _run_drop_columns(self, backwards=False):

_Table: t.Type[Table] = create_table_class(
class_name=table_class_name,
class_kwargs={"tablename": columns[0].tablename},
class_kwargs={
"tablename": columns[0].tablename,
"schema": columns[0].schema,
},
)

for column in columns:
Expand All @@ -662,7 +679,11 @@ async def _run_rename_tables(self, backwards=False):
)

_Table: t.Type[Table] = create_table_class(
class_name=class_name, class_kwargs={"tablename": tablename}
class_name=class_name,
class_kwargs={
"tablename": tablename,
"schema": rename_table.schema,
},
)

await self._run_query(
Expand All @@ -680,7 +701,10 @@ async def _run_rename_columns(self, backwards=False):

_Table: t.Type[Table] = create_table_class(
class_name=table_class_name,
class_kwargs={"tablename": columns[0].tablename},
class_kwargs={
"tablename": columns[0].tablename,
"schema": columns[0].schema,
},
)

for rename_column in columns:
Expand Down Expand Up @@ -746,7 +770,10 @@ async def _run_add_columns(self, backwards=False):

_Table: t.Type[Table] = create_table_class(
class_name=add_column.table_class_name,
class_kwargs={"tablename": add_column.tablename},
class_kwargs={
"tablename": add_column.tablename,
"schema": add_column.schema,
},
)

await self._run_query(
Expand All @@ -765,7 +792,10 @@ async def _run_add_columns(self, backwards=False):
# sets up the columns correctly.
_Table: t.Type[Table] = create_table_class(
class_name=add_columns[0].table_class_name,
class_kwargs={"tablename": add_columns[0].tablename},
class_kwargs={
"tablename": add_columns[0].tablename,
"schema": add_columns[0].schema,
},
class_members={
add_column.column._meta.name: add_column.column
for add_column in add_columns
Expand Down
5 changes: 5 additions & 0 deletions piccolo/apps/migrations/auto/operations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ class RenameTable:
old_tablename: str
new_class_name: str
new_tablename: str
schema: t.Optional[str] = None


@dataclass
Expand All @@ -28,6 +29,7 @@ class RenameColumn:
new_column_name: str
old_db_column_name: str
new_db_column_name: str
schema: t.Optional[str] = None


@dataclass
Expand All @@ -40,6 +42,7 @@ class AlterColumn:
old_params: t.Dict[str, t.Any]
column_class: t.Optional[t.Type[Column]] = None
old_column_class: t.Optional[t.Type[Column]] = None
schema: t.Optional[str] = None


@dataclass
Expand All @@ -48,6 +51,7 @@ class DropColumn:
column_name: str
db_column_name: str
tablename: str
schema: t.Optional[str] = None


@dataclass
Expand All @@ -58,3 +62,4 @@ class AddColumn:
column_class_name: str
column_class: t.Type[Column]
params: t.Dict[str, t.Any]
schema: t.Optional[str] = None
31 changes: 27 additions & 4 deletions piccolo/apps/migrations/auto/schema_differ.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,7 @@ def check_rename_tables(self) -> RenameTableCollection:
old_tablename=drop_table.tablename,
new_class_name=new_table.class_name,
new_tablename=new_table.tablename,
schema=new_table.schema,
)
)
break
Expand All @@ -212,6 +213,7 @@ def check_rename_tables(self) -> RenameTableCollection:
old_tablename=drop_table.tablename,
new_class_name=new_table.class_name,
new_tablename=new_table.tablename,
schema=new_table.schema,
)
)
break
Expand Down Expand Up @@ -288,6 +290,7 @@ def check_renamed_columns(self) -> RenameColumnCollection:
new_column_name=add_column.column_name,
old_db_column_name=drop_column.db_column_name,
new_db_column_name=add_column.db_column_name,
schema=add_column.schema,
)
)
break
Expand Down Expand Up @@ -516,8 +519,14 @@ def alter_columns(self) -> AlterStatements:
)
)

schema_str = (
"None"
if alter_column.schema is None
else f'"{alter_column.schema}"'
)

response.append(
f"manager.alter_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{alter_column.column_name}', db_column_name='{alter_column.db_column_name}', params={new_params.params}, old_params={old_params.params}, column_class={column_class}, old_column_class={old_column_class})" # noqa: E501
f"manager.alter_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{alter_column.column_name}', db_column_name='{alter_column.db_column_name}', params={new_params.params}, old_params={old_params.params}, column_class={column_class}, old_column_class={old_column_class}, schema={schema_str})" # noqa: E501
)

return AlterStatements(
Expand All @@ -543,8 +552,12 @@ def drop_columns(self) -> AlterStatements:
):
continue

schema_str = (
"None" if column.schema is None else f'"{column.schema}"'
)

response.append(
f"manager.drop_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{column.column_name}', db_column_name='{column.db_column_name}')" # noqa: E501
f"manager.drop_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{column.column_name}', db_column_name='{column.db_column_name}', schema={schema_str})" # noqa: E501
)
return AlterStatements(statements=response)

Expand Down Expand Up @@ -585,8 +598,14 @@ def add_columns(self) -> AlterStatements:
)
)

schema_str = (
"None"
if add_column.schema is None
else f'"{add_column.schema}"'
)

response.append(
f"manager.add_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{add_column.column_name}', db_column_name='{add_column.db_column_name}', column_class_name='{add_column.column_class_name}', column_class={column_class.__name__}, params={str(cleaned_params)})" # noqa: E501
f"manager.add_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{add_column.column_name}', db_column_name='{add_column.db_column_name}', column_class_name='{add_column.column_class_name}', column_class={column_class.__name__}, params={str(cleaned_params)}, schema={schema_str})" # noqa: E501
)
return AlterStatements(
statements=response,
Expand Down Expand Up @@ -647,8 +666,12 @@ def new_table_columns(self) -> AlterStatements:
)
)

schema_str = (
"None" if table.schema is None else f'"{table.schema}"'
)

response.append(
f"manager.add_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{column._meta.name}', db_column_name='{column._meta.db_column_name}', column_class_name='{column.__class__.__name__}', column_class={column.__class__.__name__}, params={str(cleaned_params)})" # noqa: E501
f"manager.add_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{column._meta.name}', db_column_name='{column._meta.db_column_name}', column_class_name='{column.__class__.__name__}', column_class={column.__class__.__name__}, params={str(cleaned_params)}, schema={schema_str})" # noqa: E501
)
return AlterStatements(
statements=response,
Expand Down
78 changes: 72 additions & 6 deletions tests/apps/migrations/auto/integration/test_migrations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -168,14 +168,14 @@ def _test_migrations(
time.sleep(1e-6)

if test_function:
column_name = (
table_snapshots[-1][-1]
._meta.non_default_columns[0]
._meta.db_column_name
)
column = table_snapshots[-1][-1]._meta.non_default_columns[0]
column_name = column._meta.db_column_name
schema = column._meta.table._meta.schema
tablename = column._meta.table._meta.tablename
row_meta = self.get_postgres_column_definition(
tablename="my_table",
tablename=tablename,
column_name=column_name,
schema=schema or "public",
)
self.assertTrue(
test_function(row_meta),
Expand Down Expand Up @@ -1213,6 +1213,72 @@ def test_move_table_to_public_schema(self):
).run_sync(),
)

def test_altering_table_in_schema(self):
"""
Make sure tables in schemas can be altered.
https://github.com/piccolo-orm/piccolo/issues/883
"""
self._test_migrations(
table_snapshots=[
# Create a table with a single column
[
create_table_class(
class_name="Manager",
class_kwargs={"schema": self.new_schema},
class_members={"first_name": Varchar()},
)
],
# Rename the column
[
create_table_class(
class_name="Manager",
class_kwargs={"schema": self.new_schema},
class_members={"name": Varchar()},
)
],
# Add a column
[
create_table_class(
class_name="Manager",
class_kwargs={"schema": self.new_schema},
class_members={
"name": Varchar(),
"age": Integer(),
},
)
],
# Remove a column
[
create_table_class(
class_name="Manager",
class_kwargs={"schema": self.new_schema},
class_members={
"name": Varchar(),
},
)
],
# Alter a column
[
create_table_class(
class_name="Manager",
class_kwargs={"schema": self.new_schema},
class_members={
"name": Varchar(length=512),
},
)
],
],
test_function=lambda x: all(
[
x.column_name == "name",
x.data_type == "character varying",
x.character_maximum_length == 512,
]
),
)


@engines_only("postgres", "cockroach")
class TestSameTableName(MigrationTestCase):
Expand Down
Loading

0 comments on commit d2b2a5a

Please sign in to comment.