Merge pull request #41 from kraken-tech/backfill-positive-negative

Set up backfill log with separate positive and negative ranges

Author: timb07

src/psycopack/_introspect.py

@@ -123,16 +123,25 @@ def get_table_columns(self, *, table: str) -> list[str]:
         return [r[0] for r in self.cur.fetchall()]
 
     def get_min_and_max_pk(
-        self, *, table: str, pk_column: str
+        self, *, table: str, pk_column: str, positive: bool
     ) -> tuple[int, int] | None:
+        if positive:
+            # positive range
+            where_sql = "WHERE {pk_column} >= 0;"
+        else:
+            # negative range
+            where_sql = "WHERE {pk_column} < 0;"
         self.cur.execute(
             psycopg.sql.SQL(
-                dedent("""
-                SELECT
-                  MIN({pk_column}) AS min_pk,
-                  MAX({pk_column}) AS max_pk
-                FROM {schema}.{table};
-                """)
+                dedent(
+                    """
+                    SELECT
+                      MIN({pk_column}) AS min_pk,
+                      MAX({pk_column}) AS max_pk
+                    FROM {schema}.{table}
+                    """
+                    + where_sql
+                )
             )
             .format(
                 table=psycopg.sql.Identifier(table),

src/psycopack/_repack.py

@@ -627,9 +627,26 @@ def _create_backfill_log(self) -> None:
         self.command.create_backfill_log(table=self.backfill_log)
 
     def _populate_backfill_log(self) -> None:
+        # positive pk values
         min_and_max = self.introspector.get_min_and_max_pk(
             table=self.table,
             pk_column=self.pk_column,
+            positive=True,
+        )
+        if min_and_max is not None:
+            min_pk, max_pk = min_and_max
+            self.command.populate_backfill_log(
+                table=self.backfill_log,
+                batch_size=self.batch_size,
+                min_pk=min_pk,
+                max_pk=max_pk,
+            )
+
+        # negative pk values
+        min_and_max = self.introspector.get_min_and_max_pk(
+            table=self.table,
+            pk_column=self.pk_column,
+            positive=False,
         )
         if min_and_max is not None:
             min_pk, max_pk = min_and_max

tests/factories.py

@@ -216,3 +216,34 @@ def create_table_for_repacking(
         SELECT generate_series(1, {referring_table_rows});
     """)
     )
+
+
+def create_table_for_backfilling(
+    cur: _cur.Cursor,
+    table_name: str = "to_backfill",
+    positive_rows: int = 100,
+    negative_rows: int = 100,
+) -> None:
+    cur.execute(
+        dedent(f"""
+            CREATE TABLE {table_name} (
+                id SERIAL PRIMARY KEY
+            );
+        """)
+    )
+
+    if positive_rows > 0:
+        cur.execute(
+            dedent(f"""
+                INSERT INTO {table_name} (id) SELECT generate_series(1, {positive_rows});
+            """)
+        )
+
+    if negative_rows > 0:
+        negative_base = -1000  # arbitrary
+        cur.execute(
+            dedent(f"""
+                INSERT INTO {table_name} (id)
+                SELECT generate_series({negative_base}, {negative_base} + {negative_rows} - 1);
+            """)
+        )

tests/test_repack.py

@@ -1989,3 +1989,40 @@ def test_when_repack_is_reinstantiated_after_swapping(
             repack=repack,
             cur=cur,
         )
+
+
+@pytest.mark.parametrize(
+    "positive_rows, negative_rows, expected_ranges",
+    [
+        (100, 0, [(1, 50), (51, 100)]),
+        (0, 100, [(-1000, -951), (-950, -901)]),
+        (100, 100, [(1, 50), (51, 100), (-1000, -951), (-950, -901)]),
+    ],
+)
+def test_populate_backfill_log(
+    connection: _psycopg.Connection,
+    positive_rows: int,
+    negative_rows: int,
+    expected_ranges: list[tuple[int, int]],
+) -> None:
+    with _cur.get_cursor(connection, logged=True) as cur:
+        factories.create_table_for_backfilling(
+            cur=cur,
+            positive_rows=positive_rows,
+            negative_rows=negative_rows,
+        )
+        repack = Psycopack(
+            table="to_backfill",
+            batch_size=50,
+            conn=connection,
+            cur=cur,
+        )
+        repack.pre_validate()
+
+        repack.setup_repacking()  # calls self._populate_backfill_log()
+
+        cur.execute(
+            f"SELECT batch_start, batch_end FROM {repack.backfill_log} ORDER BY id;"
+        )
+        result = cur.fetchall()
+        assert result == expected_ranges