#!/usr/bin/env python3
"""
migrate_users.py

Source DB -> Target DB 로 특정 테이블 데이터를 통째로 옮깁니다.
사용 예시:
    python migrate_users.py \
      --source "postgresql://postgres:pwd@oldhost:5432/postgres" \
      --target "postgresql://postgres:pwd@newhost:5432/postgres"
"""

import argparse
import sys
import psycopg2
from psycopg2 import sql

TABLES = [
    ("public", "users"),
    ("public", "user_banned_words"),
    ("public", "common_banned_words_for_kipris"),
]

def copy_table(src_conn, dst_conn, schema, table):
    """
    src_conn: psycopg2 connection to source DB
    dst_conn: psycopg2 connection to target DB
    """
    src_cur = src_conn.cursor()
    dst_cur = dst_conn.cursor()

    qualified = sql.Identifier(schema) + sql.SQL('.') + sql.Identifier(table)

    # 1) 모든 데이터를 CSV 형식으로 추출
    src_sql = sql.SQL("COPY {} TO STDOUT WITH (FORMAT csv, HEADER false)").format(qualified)
    # 2) 대상 DB로 그대로 복사
    dst_sql = sql.SQL("COPY {} FROM STDIN WITH (FORMAT csv, HEADER false)").format(qualified)

    print(f">>> Migrating {schema}.{table} ...", file=sys.stderr)
    src_cur.copy_expert(src_sql.as_string(src_conn), sys.stdout if False else None)  # 드라이런 방지
    # 대신 스트림으로 처리:
    buffer = []
    def writer(data):
        buffer.append(data)
    src_cur.copy_expert(src_sql.as_string(src_conn), writer)

    # 이제 WRITE to target
    # 주의: 이미 데이터가 있을 경우를 대비해 TRUNCATE 또는 DELETE 할지 결정하세요
    dst_cur.execute(sql.SQL("TRUNCATE {} CASCADE").format(qualified))
    dst_conn.commit()

    # 스트림 복사
    dst_cur.copy_expert(dst_sql.as_string(dst_conn), "\n".join(buffer))
    dst_conn.commit()

    print(f">>> Finished {schema}.{table}", file=sys.stderr)

def main():
    parser = argparse.ArgumentParser(description="Migrate selected tables from one Postgres to another")
    parser.add_argument("--source", "-s", required=True, help="Source DB URL")
    parser.add_argument("--target", "-t", required=True, help="Target DB URL")
    args = parser.parse_args()

    # 1) 두 커넥션 열기
    src_conn = psycopg2.connect(args.source)
    dst_conn = psycopg2.connect(args.target)

    try:
        for schema, table in TABLES:
            copy_table(src_conn, dst_conn, schema, table)
    finally:
        src_conn.close()
        dst_conn.close()

if __name__ == "__main__":
    main()



# # 예시 데이터
# chmod +x migrate_users.py
# python migrate_users.py \
#   --source "postgresql://postgres:oldpwd@127.0.0.1:54322/postgres" \
#   --target "postgresql://postgres:newpwd@127.0.0.1:54322/postgres"