Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Glue code to quickly copy data from one Postgres table to another
#!/usr/bin/env python3
import click
import queue
import sys
import os
import threading
q = queue.Queue()
EXPECTED_DEST_DB_HOST = "host=ur-db.postgres.database"
@click.command()
@click.option("--table-name", help="Name of table to migrate", required=True)
@click.option(
"--threads",
"-j",
help="Number of threads to run in parallel",
type=int,
)
@click.option(
"--table-size", help="Number of rows in the table", type=int, required=True
)
@click.option("--source-url", help="URL of the source database", required=True)
@click.option("--dest-url", help="URL of the source database", required=True)
@click.option("--dry-run/--no-dry-run", default=True, help="Is this a dry run")
@click.option("--batch-size", default=1000, type=int)
@click.option("--start", default=0, type=int, help="Lowest ID to migrate")
def migrate_table(
dest_url: str,
source_url: str,
table_name: str,
threads: int,
table_size: int,
dry_run: bool,
batch_size: int,
start: int,
) -> None:
"""Copy the specified table from the database at source_url to the database at
dest_url. The copy happens in parallel (according to the value of the --threads
param) using os.system.
"""
if EXPECTED_DEST_DB_HOST not in dest_url:
print("Error: dest database is unexpected")
sys.exit(1)
interval = batch_size
end = start + interval
while start < table_size:
task = {
"dry_run": dry_run,
"table_name": table_name,
"source_url": source_url,
"dest_url": dest_url,
"start": start,
"end": end,
}
q.put(task)
start = end
end += interval
print("starting workers")
for _ in range(threads):
threading.Thread(target=worker, daemon=True).start()
print("waiting..")
q.join()
print("done")
def worker():
print("booting a worker")
while True:
item = q.get()
dry_run = item["dry_run"]
table_name = item["table_name"]
source_url = item["source_url"]
dest_url = item["dest_url"]
start = item["start"]
end = item["end"]
if EXPECTED_DEST_DB_HOST not in dest_url:
print("Error: dest database is unexpected")
return
select_query = f"COPY (SELECT * FROM {table_name} WHERE id >= {start} AND id < {end}) TO STDOUT"
read_command = f'psql "{source_url}" -c "{select_query}"'
write_command = f'psql "{dest_url}" -c "COPY {table_name} FROM STDIN"'
command = f"PGOPTIONS='--statement-timeout=0' {read_command} | PGOPTIONS='--statement-timeout=0' {write_command}"
if dry_run:
print(f"would run:\n\t{command}")
else:
os.system(command)
print(f"migrated [{start}, {end})")
q.task_done()
if __name__ == "__main__":
migrate_table()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment