如何为Redshift Unload生成的文件添加指定表头行（Python实现）

我来帮你搞定给Redshift Unload输出文件添加指定格式表头的问题！下面分两种场景给出实用方案，你可以根据自己的实际情况选择：

方案1：直接在Redshift UNLOAD命令中添加自定义表头 #

如果你的UNLOAD查询还没最终确定，优先推荐这个方案——直接让Redshift在导出时就写入指定表头，不需要后续Python额外处理，效率更高。

只需要在UNLOAD语句里加上自定义的HEADER参数，把你需要的表头字符串直接填进去就行，示例如下：

UNLOAD ('SELECT col1, col2, col3, ... FROM your_target_table WHERE ...')
TO 's3://your-bucket/path/your_output_prefix_'
IAM_ROLE 'arn:aws:iam::123456789012:role/your-redshift-iam-role'
FORMAT AS CSV
HEADER '~id ~from ~to ~label starttime:Date endtime:Date accountid:String srcaddr:String dstaddr:String srcport:String dstport:String protocol:String packets:Int bytes:Int action:String'
DELIMITER ' ' -- 注意要和你的数据分隔符匹配，这里表头用空格分隔，所以设为空格
ALLOWOVERWRITE;

⚠️ 注意：确保DELIMITER和你实际导出的数据分隔符一致，不然表头和数据的列对应会出问题。

方案2：用Python脚本为已生成的文件补加表头 #

如果已经通过UNLOAD生成了无表头的文件，需要用Python处理每个文件添加指定表头，下面是完善后的代码示例（假设文件存在S3上，这是Redshift UNLOAD的默认输出位置）：

import datetime
import boto3
from io import StringIO

def handler(event, context):
    now = datetime.utcnow().replace(second=0, microsecond=0)
    past = now - datetime.timedelta(hours=1) # 这里按你的时间逻辑调整，比如取过去1小时的数据
    
    # 配置你的S3存储信息
    s3_bucket = 'your-bucket-name'
    s3_prefix = f'path/to/your/outputs/{past.strftime("%Y%m%d%H")}_' # 匹配UNLOAD输出的文件前缀
    
    # 定义你需要的固定表头行
    target_header = '~id ~from ~to ~label starttime:Date endtime:Date accountid:String srcaddr:String dstaddr:String srcport:String dstport:String protocol:String packets:Int bytes:Int action:String'
    
    s3_client = boto3.client('s3')
    
    # 列出所有匹配前缀的输出文件
    file_list = s3_client.list_objects_v2(Bucket=s3_bucket, Prefix=s3_prefix)
    if 'Contents' not in file_list:
        print("未找到符合条件的输出文件")
        return
    
    for file_obj in file_list['Contents']:
        file_key = file_obj['Key']
        # 跳过S3上的文件夹对象
        if file_key.endswith('/'):
            continue
        
        # 读取原文件内容
        raw_content = s3_client.get_object(Bucket=s3_bucket, Key=file_key)['Body'].read().decode('utf-8')
        # 在内容开头插入表头
        updated_content = f"{target_header}\n{raw_content}"
        
        # 写回S3覆盖原文件（也可以指定新路径避免覆盖）
        s3_client.put_object(
            Bucket=s3_bucket,
            Key=file_key,
            Body=updated_content.encode('utf-8')
        )
        
        print(f"已为文件 {file_key} 添加表头")

特殊情况处理：如果文件是GZIP压缩格式 #

如果你的UNLOAD用了GZIP压缩，需要先解压处理再重新压缩，代码调整如下：

import gzip

# 读取压缩文件
raw_obj = s3_client.get_object(Bucket=s3_bucket, Key=file_key)
with gzip.GzipFile(fileobj=raw_obj['Body'], mode='rb') as gz_file:
    raw_content = gz_file.read().decode('utf-8')

# 生成带表头的压缩内容
updated_compressed = gzip.compress(f"{target_header}\n{raw_content}".encode('utf-8'))
# 写回S3
s3_client.put_object(Bucket=s3_bucket, Key=file_key, Body=updated_compressed)

如果是本地文件处理 #

如果文件在本地磁盘，代码会更简单：

import os

def add_header_to_local_files(target_folder):
    target_header = '~id ~from ~to ~label starttime:Date endtime:Date accountid:String srcaddr:String dstaddr:String srcport:String dstport:String protocol:String packets:Int bytes:Int action:String'
    
    for filename in os.listdir(target_folder):
        file_path = os.path.join(target_folder, filename)
        if not os.path.isfile(file_path):
            continue
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(f"{target_header}\n{content}")
        
        print(f"已为本地文件 {file_path} 添加表头")

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