Add script to clone inventory between tenant DBs

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()