add tablespace argument to Table

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] = {

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,
             )
         )
 

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:

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
         ]

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
                 )
 

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]:

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:

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