优化Snowflake到PostgreSQL数据迁移效率，寻求直接迁移方案

我正在编写Python脚本，将Snowflake中的表迁移至PostgreSQL数据库，要求把Snowflake表的每行转换为JSON字符串，最终PostgreSQL表仅包含索引列和存储该行全部数据的JSON字符串。这部分转换逻辑已实现，但迁移耗时过长：当前方案是通过Snowflake的COPY命令将表导出至S3存储桶的csv.gz文件，再使用psycopg2库连接PostgreSQL，通过copy_expert方法结合COPY语句加载S3文件到目标表。迁移50-60GB数据需耗时2-3小时，想了解是否可通过Python实现Snowflake到PostgreSQL的直接数据迁移，或其他优化方案。

当前核心实现代码 #

import os
import psycopg2
import boto3
from connections import SnFlDWH

# copy table from snowflake into s3
query = f"COPY INTO 's3://{{bucket_name}}/{{s3_key}}'
FROM (
    SELECT * 
    FROM {{db_name}}.{{schema_name}}.{{table_name}}
)
FILE_FORMAT = (
    TYPE = CSV,
    FIELD_OPTIONALLY_ENCLOSED_BY = '\"' 
)
OVERWRITE = TRUE
CREDENTIALS = (
    AWS_KEY_ID='{{aws_key}}',
    AWS_SECRET_KEY='{{aws_secret}}'
);"

# establish a connection to snowflake and execute the query
snfl_conn = SnFlDWH()
snfl_conn.execqute_sql(query)

# establish a connection to postgreSQL db
pstg_conn = psycopg2.connect(
        host=host,
        database=database,
        user=user,
        password=password
    )

cursor = conn.cursor()

# connect to s3 and get a list of all files on specific s3 bucket
# download each csv.gz file and read the file in locally
# load each csv file into the specified postgreSQL table using the COPY command
s3_client = boto3.client('s3', aws_key=aws_key, aws_secret=aws_secret, region_name=region_name)
response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=s3_key)
for obj in response.get('Contents', []):
    file_name = obj['Key']
    if file_name.endswith('.csv.gz'):
        # Download files from S3
        local_file_name = f"/tmp/{os.path.basename(file_name)}"
        s3_client.download_file(bucket_name, file_name, local_file_name)
        #Load data from local file into Postgres table
        with gzip.open(local_file_name, 'rb') as f:
            cursor.copy_expert(f"COPY {{postgres_table_name}} FROM STDIN WITH CSV DELIMITER ','", f)
        # Clean up: Delete the downloaded file
        os.remove(local_file_name)
        #Delete file from S3
        s3_client.delete_object(Bucket=bucket_name, Key=file_name)
        
# Commit the transaction and close the cursor and connection
conn.commit()
cursor.close()
conn.close()

优化方案建议 #

一、直接从Snowflake流式写入PostgreSQL（Python实现） #

跳过S3中转环节，直接拉取Snowflake数据并写入PostgreSQL，减少中间IO开销：

• 使用Snowflake官方Python Connector（snowflake-connector-python）批量查询数据，设置合理的fetch_size（如10000行），避免内存过载。
• 查询时直接将每行转换为JSON字符串，结合PostgreSQL的批量写入方法提升效率。示例逻辑：

  import snowflake.connector
  import psycopg2
  import json
  
  # Snowflake连接配置
  sf_conn = snowflake.connector.connect(
      user=sf_user,
      password=sf_password,
      account=sf_account,
      warehouse=sf_warehouse,
      database=sf_db,
      schema=sf_schema
  )
  sf_cursor = sf_conn.cursor()
  sf_cursor.execute("SELECT * FROM your_snowflake_table")
  
  # PostgreSQL连接配置
  pg_conn = psycopg2.connect(
      host=pg_host,
      database=pg_db,
      user=pg_user,
      password=pg_password
  )
  pg_cursor = pg_conn.cursor()
  
  # 批量拉取并写入
  batch_size = 10000
  while True:
      rows = sf_cursor.fetchmany(batch_size)
      if not rows:
          break
      # 转换为JSON字符串，适配PostgreSQL表结构（假设仅需JSON字段）
      values = [(json.dumps(row),) for row in rows]
      pg_cursor.executemany("INSERT INTO your_pg_table (data_json) VALUES (%s)", values)
      pg_conn.commit()
  
  # 关闭连接
  sf_cursor.close()
  sf_conn.close()
  pg_cursor.close()
  pg_conn.close()
  

• 注意：迁移前可临时禁用PostgreSQL表的索引和约束，完成后重建，大幅提升写入速度。

二、优化现有S3中转方案 #

若保留S3中转，可从以下维度优化：

1. Snowflake导出优化
   • 直接导出JSON格式文件，避免CSV转JSON的额外开销。修改COPY语句：

     COPY INTO 's3://bucket/path/'
     FROM (SELECT OBJECT_CONSTRUCT(*) FROM your_snowflake_table)
     FILE_FORMAT = (TYPE = JSON, COMPRESSION = GZIP)
     OVERWRITE = TRUE
     CREDENTIALS = (AWS_KEY_ID='xxx', AWS_SECRET_KEY='xxx')
     MAX_FILE_SIZE = 5368709120; -- 设置为5GB，减少小文件数量
     

   • 临时扩容Snowflake Warehouse，提升导出速度，完成后恢复原规格。
2. PostgreSQL加载优化
   • 使用PostgreSQL官方aws_s3扩展，直接从S3加载数据，无需下载到本地：

     CREATE EXTENSION IF NOT EXISTS aws_s3;
     COPY your_pg_table (data_json)
     FROM S3 's3://bucket/path/'
     WITH (FORMAT 'json', COMPRESSION 'gzip', AWS_KEY_ID 'xxx', AWS_SECRET_KEY 'xxx');
     

   • 迁移前执行SET synchronous_commit = off;，降低写入等待时间；关闭自动提交，批量提交事务。
3. 并行处理
   • 用多线程或多进程同时处理多个S3文件的加载任务，充分利用CPU和网络资源。

三、通用环境优化 #

• 确保Snowflake、S3、PostgreSQL部署在同一云服务商区域，减少跨区域数据传输延迟。
• 临时调优PostgreSQL配置：增大work_mem、wal_buffers参数，提升批量写入性能，迁移完成后恢复默认值。

内容的提问来源于stack exchange，提问作者Olek