added get_table_from_snaphot to MigrationManager, along with offset

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -322,6 +322,40 @@ def deserialise_params(
 
     ###########################################################################
 
+    def get_table_from_snaphot(
+        self,
+        app_name: str,
+        table_class_name: str,
+        offset: int = 0,
+        migration_id: t.Optional[str] = None,
+    ) -> t.Type[Table]:
+        """
+        Returns a Table subclass which can be used for modifying data within
+        a migration.
+
+        :param offset:
+            Lets you get a table as it appeared in an older migration. If the
+            offset is -1, the table will come from the previous migration.
+
+        """
+        from piccolo.apps.migrations.commands.base import BaseMigrationManager
+
+        if migration_id is None:
+            migration_id = self.migration_id
+
+        return (
+            BaseMigrationManager()
+            .get_table_from_snaphot(
+                app_name=app_name,
+                table_class_name=table_class_name,
+                max_migration_id=migration_id,
+                offset=offset,
+            )
+            .to_table_class()
+        )
+
+    ###########################################################################
+
     async def _run_alter_columns(self, backwards=False):
         for table_class_name in self.alter_columns.table_class_names:
             alter_columns = self.alter_columns.for_table_class_name(

piccolo/apps/migrations/commands/base.py

@@ -126,7 +126,10 @@ def get_migration_ids(
         return sorted(list(migration_module_dict.keys()))
 
     def get_migration_managers(
-        self, app_name: str, max_migration_id: t.Optional[str] = None
+        self,
+        app_name: str,
+        max_migration_id: t.Optional[str] = None,
+        offset: int = 0,
     ) -> t.List[MigrationManager]:
         """
         :param max_migration_id:
@@ -154,20 +157,28 @@ def get_migration_managers(
                 else:
                     migration_managers.append(response)
 
-        return migration_managers
+        if offset > 0:
+            raise Exception(
+                "Positive offset values aren't currently supported"
+            )
+        elif offset < 0:
+            return migration_managers[0:offset]
+        else:
+            return migration_managers
 
     def get_table_from_snaphot(
         self,
         app_name: str,
         table_class_name: str,
         max_migration_id: t.Optional[str] = None,
+        offset: int = 0,
     ) -> DiffableTable:
         """
         This will generate a SchemaSnapshot up to the given migration ID, and
         will return a DiffableTable class from that snapshot.
         """
         migration_managers = self.get_migration_managers(
-            app_name=app_name, max_migration_id=max_migration_id
+            app_name=app_name, max_migration_id=max_migration_id, offset=offset
         )
         schema_snapshot = SchemaSnapshot(managers=migration_managers)