AutoPercenty3/test/supabase_migration/main.py

85 lines
2.6 KiB
Python

#!/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"