如何在Apache Beam Java中从.sql文件传入SQL语句

解决方案 #

Java版Apache Beam没有内置直接读取SQL文件并传入SpannerIO.read().withQuery()的API，但可以通过普通Java IO或Beam FileIO实现，分场景处理：

1. 本地开发/测试场景 #

如果是本地运行作业，直接用Java标准IO读取本地SQL文件（Java 11+推荐Files.readString）：

import java.nio.file.Files;
import java.nio.file.Paths;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
import org.apache.beam.sdk.values.PCollection;
import com.google.cloud.spanner.Struct;

public class SpannerSqlFromFile {
    public static void main(String[] args) {
        Pipeline pipeline = Pipeline.create();
        
        // 读取本地SQL文件内容
        String spnQuery;
        try {
            spnQuery = Files.readString(Paths.get("./queries/my-spanner-query.sql"));
        } catch (Exception e) {
            throw new RuntimeException("Failed to read SQL file", e);
        }
        
        // 执行Spanner查询
        PCollection<Struct> queryResults = pipeline.apply(
            SpannerIO.read()
                .withInstanceId("your-spanner-instance-id")
                .withDatabaseId("your-spanner-db-id")
                .withQuery(spnQuery)
        );
        
        // 后续处理查询结果...
        
        pipeline.run().waitUntilFinish();
    }
}

2. 分布式运行场景（如Dataflow） #

如果作业要在分布式环境运行，SQL文件需存放在GCS/S3等分布式存储，用Beam FileIO读取：

import org.apache.beam.sdk.io.FileIO;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;
import com.google.cloud.spanner.Struct;

public class DistributedSpannerSql {
    public static void main(String[] args) {
        Pipeline pipeline = Pipeline.create();
        
        // 读取GCS上的SQL文件
        PCollection<String> sqlContent = pipeline.apply(
                FileIO.match().filepattern("gs://your-bucket/path/to/query.sql")
            )
            .apply(FileIO.readMatches())
            .apply(ParDo.of(new DoFn<FileIO.ReadableFile, String>() {
                @ProcessElement
                public void processElement(ProcessContext c) throws Exception {
                    c.output(c.element().readFullyAsUTF8String());
                }
            }));
        
        // 将多元素集合转为单元素（确保SQL文件唯一）
        PCollection<String> singleQuery = sqlContent.apply(
            Combine.globally(inputs -> {
                var iter = inputs.iterator();
                return iter.hasNext() ? iter.next() : "";
            }).withoutDefaults()
        );
        
        // 传递查询语句执行Spanner查询
        singleQuery.apply(ParDo.of(new DoFn<String, Struct>() {
            @ProcessElement
            public void processElement(ProcessContext c) {
                String query = c.element();
                // 这里可以嵌套SpannerIO读取，或通过侧输出传递给其他步骤
                PCollection<Struct> results = c.getPipeline().apply(
                    SpannerIO.read()
                        .withInstanceId("instance-id")
                        .withDatabaseId("db-id")
                        .withQuery(query)
                );
                // 处理results...
            }
        }));
        
        pipeline.run().waitUntilFinish();
    }
}

3. 推荐方案：通过PipelineOptions传递 #

更简洁的方式是在作业启动前读取SQL文件，将内容存入自定义PipelineOptions，避免分布式环境的文件读取问题：

import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import java.nio.file.Files;
import java.nio.file.Paths;

// 自定义PipelineOptions
public interface SpannerSqlOptions extends PipelineOptions {
    @Description("Path to SQL query file (local or GCS path)")
    @Default.String("./query.sql")
    String getSqlFilePath();
    void setSqlFilePath(String value);
}

public class SpannerSqlWithOptions {
    public static void main(String[] args) {
        SpannerSqlOptions options = PipelineOptionsFactory
            .fromArgs(args)
            .withValidation()
            .as(SpannerSqlOptions.class);
        
        // 读取SQL文件内容（本地路径直接读，GCS路径可通过GCS API读取）
        String spnQuery;
        try {
            if (options.getSqlFilePath().startsWith("gs://")) {
                // 处理GCS路径：用Google Cloud Storage API直接读取
                String bucketAndPath = options.getSqlFilePath().replace("gs://", "");
                int splitIndex = bucketAndPath.indexOf('/');
                String bucketName = bucketAndPath.substring(0, splitIndex);
                String objectPath = bucketAndPath.substring(splitIndex + 1);
                
                spnQuery = com.google.cloud.storage.StorageOptions.getDefaultInstance()
                    .getService()
                    .get(bucketName, objectPath)
                    .getContent()
                    .toString();
            } else {
                spnQuery = Files.readString(Paths.get(options.getSqlFilePath()));
            }
        } catch (Exception e) {
            throw new RuntimeException("Failed to read SQL file", e);
        }
        
        Pipeline pipeline = Pipeline.create(options);
        pipeline.apply(SpannerIO.read()
            .withInstanceId("instance-id")
            .withDatabaseId("db-id")
            .withQuery(spnQuery));
        
        pipeline.run().waitUntilFinish();
    }
}

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