diff --git a/scripts/clone_inventory.py b/scripts/clone_inventory.py new file mode 100644 index 0000000..fc0d84f --- /dev/null +++ b/scripts/clone_inventory.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Clone inventory tables from one tenant DB to another. + +Uses COPY with the columns common to both source and target, so minor schema +mismatches (e.g. missing latitude/longitude in branches) are handled. + +Example: + python scripts/clone_inventory.py \ + --source tenant_autopartes_estrada \ + --target tenant_originales_autopartes +""" + +import argparse +import sys +import tempfile +import time +from contextlib import closing + +import psycopg2 + + +# (table, source_where_clause) +TABLES = [ + ("branches", "id <> 1"), + ("inventory", None), + ("inventory_stock", None), + ("inventory_sku_aliases", None), + ("inventory_stock_summary", None), + ("inventory_vehicle_compat", None), +] + +SEQUENCES = [ + ("branches_id_seq", "branches"), + ("inventory_id_seq", "inventory"), + ("inventory_stock_id_seq", "inventory_stock"), + ("inventory_sku_aliases_id_seq", "inventory_sku_aliases"), + ("inventory_vehicle_compat_id_seq", "inventory_vehicle_compat"), +] + + +def connect(db_name: str): + return psycopg2.connect(host="localhost", user="postgres", dbname=db_name) + + +def get_columns(cur, table: str): + cur.execute( + """ + SELECT column_name + FROM information_schema.columns + WHERE table_name = %s + AND table_schema = 'public' + ORDER BY ordinal_position + """, + (table,), + ) + return [r[0] for r in cur.fetchall()] + + +def copy_table(src_conn, dst_conn, table: str, where: str | None): + with src_conn.cursor() as src_cur, dst_conn.cursor() as dst_cur: + src_cols = get_columns(src_cur, table) + dst_cols = set(get_columns(dst_cur, table)) + common = [c for c in src_cols if c in dst_cols] + if not common: + print(f"Skipping {table}: no common columns") + return + + col_sql = ", ".join(f'"{c}"' for c in common) + copy_to = f'COPY (SELECT {col_sql} FROM "{table}"' + if where: + copy_to += f" WHERE {where}" + copy_to += ") TO STDOUT" + copy_from = f'COPY "{table}" ({col_sql}) FROM STDIN' + + print(f"Copying {table} ({len(common)} columns)...", end=" ", flush=True) + start = time.time() + + with tempfile.SpooledTemporaryFile(max_size=50 * 1024 * 1024, mode="w+b") as tmp: + src_cur.copy_expert(copy_to, tmp) + tmp.seek(0) + dst_cur.copy_expert(copy_from, tmp) + + dst_conn.commit() + elapsed = time.time() - start + print(f"done in {elapsed:.1f}s") + + +def reset_sequences(dst_conn): + with dst_conn.cursor() as cur: + for seq, table in SEQUENCES: + cur.execute( + f"SELECT setval('{seq}', COALESCE((SELECT MAX(id) FROM \"{table}\"), 1), true)" + ) + dst_conn.commit() + + +def main(): + parser = argparse.ArgumentParser(description="Clone inventory between tenant DBs") + parser.add_argument("--source", required=True) + parser.add_argument("--target", required=True) + args = parser.parse_args() + + src = connect(args.source) + dst = connect(args.target) + + try: + # Prepare target: remove extra branches, truncate inventory tables. + with dst.cursor() as cur: + print("Preparing target tables...", end=" ", flush=True) + cur.execute("DELETE FROM branches WHERE id <> 1") + cur.execute( + """ + TRUNCATE TABLE inventory, + inventory_stock, + inventory_sku_aliases, + inventory_stock_summary, + inventory_vehicle_compat + CASCADE + """ + ) + dst.commit() + print("done") + + for table, where in TABLES: + copy_table(src, dst, table, where) + + print("Resetting sequences...", end=" ", flush=True) + reset_sequences(dst) + print("done") + + print("Inventory clone completed.") + except Exception as e: + print(f"ERROR: {e}", file=sys.stderr) + dst.rollback() + raise + finally: + src.close() + dst.close() + + +if __name__ == "__main__": + main()