从PostgreSQL读取后用TextIO.write()写入GCS如何动态设置CSV表头？

由于TextIO.Write.withHeader()仅支持静态字符串参数，无法直接通过ValueProvider传递运行时表头，以下提供三种适配不同场景的解决方案：

场景1：无需构建Dataflow模板（直接运行管道） #

如果你的管道不需要打包成模板，可在管道构建阶段直接解析ValueProvider的值，传入withHeader()：

// 从PipelineOptions中获取表头的ValueProvider
ValueProvider<String> headerProvider = options.getHeader();
// 直接解析值（仅非模板运行时可用，构建模板时调用get()会报错）
String header = headerProvider.get();

output.apply(
    "Write CSV",
    TextIO.write()
        .to(options.getFilePath())
        .withSuffix(".csv")
        .withHeader(header)
        .withShardNameTemplate("-S-of-N")
        .withTempDirectory(options.getTempDirectory())
);

场景2：需要构建Dataflow模板（运行时动态传递表头） #

这种场景下无法在构建阶段解析ValueProvider，需通过自定义写入逻辑实现动态表头：

方法：使用FileIO自定义Writer控制表头写入 #

通过FileIO.write()配合自定义FileWrite，确保每个输出文件仅写入一次表头，即使数据为空也会生成表头：

// 从PipelineOptions获取表头的ValueProvider
ValueProvider<String> headerProvider = options.getHeader();

output.apply(
    "Write CSV with Dynamic Header",
    FileIO.<String>write()
        .to(options.getFilePath())
        .withSuffix(".csv")
        .withShardNameTemplate("-S-of-N")
        .withTempDirectory(options.getTempDirectory())
        .via(TextIO.sink()) // 复用TextIO的文本写入逻辑
        .withDynamicDestinations(new DynamicDestinations<String, Void>() {
            @Override
            public Void getDestination(String element) {
                // 所有元素写入同一目标（如需分区可自定义逻辑）
                return null;
            }

            @Override
            public FileWrite<String> getWriter(Void destination) {
                return new FileWrite<String>() {
                    private boolean headerWritten = false;
                    private final TextSink.TextWriter textWriter = TextIO.sink().createWriter();

                    @Override
                    public void write(String element) throws IOException {
                        if (!headerWritten) {
                            // 运行时解析表头并写入
                            String header = headerProvider.get();
                            textWriter.write(header);
                            headerWritten = true;
                        }
                        textWriter.write(element);
                    }

                    @Override
                    public void flush() throws IOException {
                        textWriter.flush();
                    }

                    @Override
                    public void close() throws IOException {
                        // 若数据为空，仍写入表头
                        if (!headerWritten) {
                            textWriter.write(headerProvider.get());
                        }
                        textWriter.close();
                    }
                };
            }
        })
);

场景3：表头来自SQL查询的列名 #

如果表头对应SQL表的列名，可从JDBC查询的元数据中动态生成表头：

// 初始化线程安全的变量存储表头
AtomicReference<String> headerRef = new AtomicReference<>();
AtomicBoolean headerSet = new AtomicBoolean(false);

// 读取JDBC数据并生成表头
PCollection<String> data = pipeline.apply(
    "Read from SQL",
    JdbcIO.read()
        .withDataSourceConfiguration(JdbcIO.DataSourceConfiguration.create(
            options.getJdbcDriver(),
            options.getJdbcUrl()
        ).withUsername(options.getDbUser()).withPassword(options.getDbPassword()))
        .withQuery(options.getQuery())
        .withRowMapper((ResultSet resultSet) -> {
            // 首次读取时获取列名生成表头
            if (!headerSet.get()) {
                ResultSetMetaData metaData = resultSet.getMetaData();
                StringBuilder header = new StringBuilder();
                for (int i = 1; i <= metaData.getColumnCount(); i++) {
                    if (i > 1) header.append(",");
                    header.append(metaData.getColumnName(i));
                }
                headerRef.set(header.toString());
                headerSet.set(true);
            }
            // 转换行数据为CSV格式
            StringBuilder row = new StringBuilder();
            for (int i = 1; i <= metaData.getColumnCount(); i++) {
                if (i > 1) row.append(",");
                row.append(resultSet.getString(i));
            }
            return row.toString();
        })
);

// 生成表头的单元素PCollection
PCollection<String> headerPC = pipeline.apply(Create.of(headerRef.get()));
// 合并表头和数据
PCollection<String> combinedData = PCollectionList.of(headerPC).and(data).apply(Flatten.pCollections());

// 写入CSV（无需withHeader）
combinedData.apply(
    "Write CSV",
    TextIO.write()
        .to(options.getFilePath())
        .withSuffix(".csv")
        .withShardNameTemplate("-S-of-N")
        .withTempDirectory(options.getTempDirectory())
);

注意：此方法若数据为空则不会写入表头；若管道分片，每个分片都会包含表头，如需避免需结合场景2的自定义Writer逻辑。

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