143 lines
4.1 KiB
Python
143 lines
4.1 KiB
Python
#!/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()
|