risingwave_connector_codec/common/protobuf/
compiler.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16
17use prost_types::FileDescriptorSet;
18use protox::Error;
19use protox::file::{ChainFileResolver, File, FileResolver, GoogleFileResolver};
20
21// name -> content
22pub fn compile_pb(
23    main_file: (String, String),
24    dependencies: impl IntoIterator<Item = (String, String)>,
25) -> Result<FileDescriptorSet, Error> {
26    struct MyResolver {
27        map: HashMap<String, String>,
28    }
29
30    impl MyResolver {
31        fn new(
32            main_file: (String, String),
33            dependencies: impl IntoIterator<Item = (String, String)>,
34        ) -> Self {
35            let map = std::iter::once(main_file).chain(dependencies).collect();
36
37            Self { map }
38        }
39    }
40
41    impl FileResolver for MyResolver {
42        fn open_file(&self, name: &str) -> Result<File, Error> {
43            if let Some(content) = self.map.get(name) {
44                Ok(File::from_source(name, content)?)
45            } else {
46                Err(Error::file_not_found(name))
47            }
48        }
49    }
50
51    // Resolves imports from a static (import path, source) table — the bundled
52    // sets Confluent treats as ambient.
53    struct BundledFileResolver(&'static [(&'static str, &'static str)]);
54
55    impl FileResolver for BundledFileResolver {
56        fn open_file(&self, name: &str) -> Result<File, Error> {
57            let found = self.0.iter().find(|(path, _)| *path == name);
58            let (_, content) = found.ok_or_else(|| Error::file_not_found(name))?;
59            File::from_source(name, content)
60        }
61    }
62
63    let main_file_name = main_file.0.clone();
64
65    // Confluent's Java / Go SDKs prioritize user overrides.
66    let mut resolver = ChainFileResolver::new();
67    resolver.add(MyResolver::new(main_file, dependencies));
68    resolver.add(BundledFileResolver(proto_src_google_type::FILES));
69    resolver.add(BundledFileResolver(proto_src_confluent::FILES));
70    resolver.add(GoogleFileResolver::new());
71
72    let fd = protox::Compiler::with_file_resolver(resolver)
73        .include_imports(true)
74        .open_file(&main_file_name)?
75        .file_descriptor_set();
76
77    Ok(fd)
78}
risingwave_connector_codec/common/protobuf/compiler.rs

risingwave_connector_codec/common/protobuf/
compiler.rs
