Skip to content

Commit 5b92e77

Browse files
jimchampjimchamp
authored
Updates to store_index backfill script (#13521)
* Add shebang to `store_index` backfill script * Use faster query to find upper bound of IDs * Add option to begin updates starting from an arbitrary `id` * Update log statement The lower bound `id` of the next batch of updates is now included in the log statements. Output buffer is now flushed every time `print` is called. * Terminate script when WAL directory has grown too large. * Update threshold logic * Disable SIM102 in backfill script
1 parent 98ddd8d commit 5b92e77

1 file changed

Lines changed: 32 additions & 4 deletions

File tree

‎scripts/backfill_store_index.py‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
#!/usr/bin/env python3
12
"""
23
Backfills new ID column in the `store_index` table.
34
"""
@@ -13,6 +14,8 @@
1314

1415
DEFAULT_CONFIG_PATH = "conf/openlibrary.yml"
1516
DEFAULT_BATCH_SIZE = 20_000
17+
DEFAULT_LOWER_BOUND = 0
18+
DEFAULT_WAL_DIR_MAX_SIZE = 0
1619

1720

1821
def init(conf_path):
@@ -25,9 +28,20 @@ def find_upper_bound():
2528
oldb = db.get_db()
2629

2730
query = """
28-
SELECT MIN(new_id) as min from store_index
31+
SELECT MAX(id) as ubound from store_index WHERE new_id IS NULL
2932
"""
30-
return next(iter(oldb.query(query)))["min"]
33+
return next(iter(oldb.query(query)))["ubound"]
34+
35+
36+
def get_wal_dir_size():
37+
oldb = db.get_db()
38+
39+
query = """
40+
SELECT count(*) AS file_count,
41+
sum(size) AS total_bytes
42+
FROM pg_ls_waldir();
43+
"""
44+
return next(iter(oldb.query(query)))["total_bytes"]
3145

3246

3347
def backfill_rows(lower_bound, upper_bound):
@@ -47,19 +61,33 @@ def main(args):
4761
max_upper_bound = find_upper_bound()
4862

4963
# Backfill new IDs in batches
50-
lower_bound = 0
64+
lower_bound = args.lower_bound
65+
iterations = 0
5166
while lower_bound < max_upper_bound and not was_shutdown_requested():
67+
if iterations % 10 == 0 and args.wal_threshold: # noqa: SIM102
68+
if args.wal_threshold * (1024**3) < get_wal_dir_size():
69+
print("WAL directory has grown larger than threshold. Stopping script.", flush=True)
70+
break
5271
start = time.perf_counter()
5372
backfill_rows(lower_bound, lower_bound + args.batch_size)
5473
lower_bound += args.batch_size
5574
elapsed = time.perf_counter() - start
56-
print(f"Chunk updated in {elapsed:.6f} seconds")
75+
print(f"Block updated in {elapsed:.6f} seconds. Next block starts with ID {lower_bound}", flush=True)
76+
iterations += 1
5777

5878

5979
def _parse_args():
6080
p = argparse.ArgumentParser(description=__doc__)
6181
p.add_argument("-c", "--config", default=DEFAULT_CONFIG_PATH, help="Path to openlibrary configuration yaml")
6282
p.add_argument("-b", "--batch-size", default=DEFAULT_BATCH_SIZE, type=int)
83+
p.add_argument("-l", "--lower-bound", default=DEFAULT_LOWER_BOUND, type=int)
84+
p.add_argument(
85+
"-t",
86+
"--wal-threshold",
87+
default=DEFAULT_WAL_DIR_MAX_SIZE,
88+
type=float,
89+
help="If non-zero, the script will stop if the WAL directory grows beyond this many GB",
90+
)
6391
p.set_defaults(func=main)
6492
return p.parse_args()
6593

0 commit comments

Comments
 (0)