diff --git a/piccolo/apps/migrations/auto/diffable_table.py b/piccolo/apps/migrations/auto/diffable_table.py index 18a7c2809..83fb6b66f 100644 --- a/piccolo/apps/migrations/auto/diffable_table.py +++ b/piccolo/apps/migrations/auto/diffable_table.py @@ -92,6 +92,7 @@ class DiffableTable: tablename: str columns: t.List[Column] = field(default_factory=list) previous_class_name: t.Optional[str] = None + tablespace: t.Optional[str] = None def __post_init__(self): self.columns_map: t.Dict[str, Column] = { diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index a03407efb..c978a5d99 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -154,13 +154,17 @@ def add_table( class_name: str, tablename: str, columns: t.Optional[t.List[Column]] = None, + tablespace: t.Optional[str] = None, ): if not columns: columns = [] self.add_tables.append( DiffableTable( - class_name=class_name, tablename=tablename, columns=columns + class_name=class_name, + tablename=tablename, + columns=columns, + tablespace=tablespace, ) ) diff --git a/piccolo/apps/migrations/auto/schema_differ.py b/piccolo/apps/migrations/auto/schema_differ.py index 2427c44a5..c57cedf18 100644 --- a/piccolo/apps/migrations/auto/schema_differ.py +++ b/piccolo/apps/migrations/auto/schema_differ.py @@ -245,12 +245,14 @@ def create_tables(self) -> AlterStatements: not in self.rename_tables_collection.new_class_names ] - return AlterStatements( - statements=[ - f"manager.add_table('{i.class_name}', tablename='{i.tablename}')" # noqa: E501 - for i in new_tables - ] - ) + statements = [] + for i in new_tables: + tablespace = f"'{i.tablespace}'" if i.tablespace else "None" + statements.append( + f"manager.add_table('{i.class_name}', tablename='{i.tablename}', tablespace={tablespace})" # noqa: E501 + ) + + return AlterStatements(statements=statements) @property def drop_tables(self) -> AlterStatements: diff --git a/piccolo/apps/migrations/commands/new.py b/piccolo/apps/migrations/commands/new.py index 3cc04fbb0..58b805dc2 100644 --- a/piccolo/apps/migrations/commands/new.py +++ b/piccolo/apps/migrations/commands/new.py @@ -169,6 +169,7 @@ async def get_alter_statements( class_name=i.__name__, tablename=i._meta.tablename, columns=i._meta.non_default_columns, + tablespace=i._meta.tablespace, ) for i in app_config.table_classes ] diff --git a/piccolo/query/methods/create.py b/piccolo/query/methods/create.py index 94292de51..d5d76ba05 100644 --- a/piccolo/query/methods/create.py +++ b/piccolo/query/methods/create.py @@ -41,6 +41,8 @@ def default_ddl(self) -> t.Sequence[str]: base = f"{prefix} {self.table._meta.tablename}" columns_sql = ", ".join(i.ddl for i in columns) create_table_ddl = f"{base} ({columns_sql})" + if self.table._meta.tablespace and self.engine_type == "postgres": + create_table_ddl += f" TABLESPACE {self.table._meta.tablespace}" create_indexes: t.List[str] = [] for column in columns: @@ -51,6 +53,7 @@ def default_ddl(self) -> t.Sequence[str]: columns=[column], method=column._meta.index_method, if_not_exists=self.if_not_exists, + tablespace=self.table._meta.tablespace, ).ddl ) diff --git a/piccolo/query/methods/create_index.py b/piccolo/query/methods/create_index.py index 594fd63b6..f42639be4 100644 --- a/piccolo/query/methods/create_index.py +++ b/piccolo/query/methods/create_index.py @@ -17,11 +17,13 @@ def __init__( columns: t.List[t.Union[Column, str]], method: IndexMethod = IndexMethod.btree, if_not_exists: bool = False, + tablespace: t.Optional[str] = None, **kwargs, ): self.columns = columns self.method = method self.if_not_exists = if_not_exists + self.tablespace = tablespace super().__init__(table, **kwargs) @property @@ -45,12 +47,13 @@ def postgres_ddl(self) -> t.Sequence[str]: tablename = self.table._meta.tablename method_name = self.method.value column_names_str = ", ".join([f'"{i}"' for i in self.column_names]) - return [ - ( - f"{self.prefix} {index_name} ON {tablename} USING " - f"{method_name} ({column_names_str})" - ) - ] + ddl = ( + f"{self.prefix} {index_name} ON {tablename} USING " + f"{method_name} ({column_names_str})" + ) + if self.tablespace: + ddl += f" TABLESPACE {self.tablespace}" + return [ddl] @property def sqlite_ddl(self) -> t.Sequence[str]: diff --git a/piccolo/table.py b/piccolo/table.py index 9ff261023..d74c7af35 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -74,6 +74,7 @@ class TableMeta: help_text: t.Optional[str] = None _db: t.Optional[Engine] = None m2m_relationships: t.List[M2M] = field(default_factory=list) + tablespace: t.Optional[str] = None # Records reverse foreign key relationships - i.e. when the current table # is the target of a foreign key. Used by external libraries such as @@ -152,6 +153,7 @@ def __init_subclass__( db: t.Optional[Engine] = None, tags: t.List[str] = None, help_text: t.Optional[str] = None, + tablespace: t.Optional[str] = None, ): # sourcery no-metrics """ Automatically populate the _meta, which includes the tablename, and @@ -170,6 +172,10 @@ def __init_subclass__( A user friendly description of what the table is used for. It isn't used in the database, but will be used by tools such a Piccolo Admin for tooltips. + :param tablespace: + Which tablespace to create the table in. This is an advanced + feature that most users won't need, and is only applicable to + Postgres. """ if tags is None: @@ -255,6 +261,7 @@ def __init_subclass__( help_text=help_text, _db=db, m2m_relationships=m2m_relationships, + tablespace=tablespace, ) for foreign_key_column in foreign_key_columns: diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index 6fccff2e5..ed0d6288f 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -32,7 +32,7 @@ def test_add_table(self): self.assertTrue(len(create_tables.statements) == 1) self.assertEqual( create_tables.statements[0], - "manager.add_table('Band', tablename='band')", + "manager.add_table('Band', tablename='band', tablespace=None)", ) new_table_columns = schema_differ.new_table_columns