use std::{
path::{Path, PathBuf},
process::Command,
};
use buffa::Message;
use buffa_descriptor::generated::descriptor::FileDescriptorSet;
use buffa_reflect::{DescriptorPool, Kind};
use crate::{
codegen::{
EventDecl, FieldDecl, render_arrow_schema, render_builders, render_lints, render_schemas,
},
lints::LintProtoType,
options::{CodegenError, read_event_options, read_field_options},
};
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub enum DescriptorSource {
#[default]
Protoc,
Precompiled(PathBuf),
}
#[derive(Debug, Default)]
pub struct Config {
files: Vec<PathBuf>,
includes: Vec<PathBuf>,
out_dir: Option<PathBuf>,
descriptor_source: DescriptorSource,
event_prefix: Option<String>,
arrow_schema: bool,
json_render: bool,
payload_scrub: bool,
otel_attribute_view: bool,
}
impl Config {
#[must_use]
pub fn new() -> Self {
Self {
arrow_schema: true,
json_render: false,
payload_scrub: true,
otel_attribute_view: true,
..Self::default()
}
}
#[must_use]
pub fn files(mut self, files: &[impl AsRef<Path>]) -> Self {
self.files
.extend(files.iter().map(|p| p.as_ref().to_owned()));
self
}
#[must_use]
pub fn include(mut self, dir: impl AsRef<Path>) -> Self {
self.includes.push(dir.as_ref().to_owned());
self
}
#[must_use]
pub fn out_dir(mut self, dir: impl AsRef<Path>) -> Self {
self.out_dir = Some(dir.as_ref().to_owned());
self
}
#[must_use]
pub fn include_obs_options(mut self) -> Self {
self.includes
.push(PathBuf::from("__obs_build_embedded_options__"));
self
}
#[must_use]
pub fn descriptor_source(mut self, src: DescriptorSource) -> Self {
self.descriptor_source = src;
self
}
#[must_use]
pub fn event_prefix(mut self, prefix: impl Into<String>) -> Self {
self.event_prefix = Some(prefix.into());
self
}
#[must_use]
pub fn with_arrow_schema(mut self, on: bool) -> Self {
self.arrow_schema = on;
self
}
#[must_use]
pub fn with_json_render(mut self, on: bool) -> Self {
self.json_render = on;
self
}
#[must_use]
pub fn with_payload_scrub(mut self, on: bool) -> Self {
self.payload_scrub = on;
self
}
#[must_use]
pub fn with_otel_attribute_view(mut self, on: bool) -> Self {
self.otel_attribute_view = on;
self
}
pub fn compile(self) -> Result<(), CodegenError> {
let out_dir = self
.out_dir
.clone()
.or_else(|| std::env::var("OUT_DIR").ok().map(PathBuf::from))
.ok_or_else(|| CodegenError::Protoc("OUT_DIR not set".into()))?;
std::fs::create_dir_all(out_dir.join("obs")).map_err(CodegenError::OutputIo)?;
let mut effective_includes = self.includes.clone();
if effective_includes
.iter()
.any(|p| p.as_os_str() == "__obs_build_embedded_options__")
{
let embed_dir = out_dir.join("obs").join("include");
materialise_embedded_options(&embed_dir).map_err(CodegenError::OutputIo)?;
effective_includes.retain(|p| p.as_os_str() != "__obs_build_embedded_options__");
effective_includes.push(embed_dir);
}
if !self.files.is_empty() {
self.invoke_buffa_build(&out_dir, &effective_includes)?;
}
let fds_bytes = self.produce_fds(&out_dir, &effective_includes)?;
let fds = FileDescriptorSet::decode_from_slice(&fds_bytes)
.map_err(|e| CodegenError::DescriptorDecode(e.to_string()))?;
let pool = DescriptorPool::from_file_descriptor_set(fds)
.map_err(|e| CodegenError::DescriptorDecode(e.to_string()))?;
let events = collect_event_decls(&pool)?;
let event_prefix = self
.event_prefix
.clone()
.or_else(|| std::env::var("OBS_EVENT_PREFIX").ok())
.unwrap_or_else(|| "Obs".to_string());
std::fs::write(
out_dir.join("obs").join("schemas.rs"),
render_schemas(&events),
)
.map_err(CodegenError::OutputIo)?;
std::fs::write(
out_dir.join("obs").join("builders.rs"),
render_builders(&events),
)
.map_err(CodegenError::OutputIo)?;
std::fs::write(
out_dir.join("obs").join("lints.rs"),
render_lints(&events, &event_prefix),
)
.map_err(CodegenError::OutputIo)?;
if self.arrow_schema {
std::fs::write(
out_dir.join("obs").join("arrow_schema.rs"),
render_arrow_schema(&events),
)
.map_err(CodegenError::OutputIo)?;
} else {
std::fs::write(
out_dir.join("obs").join("arrow_schema.rs"),
"// arrow_schema disabled by `with_arrow_schema(false)`\n",
)
.map_err(CodegenError::OutputIo)?;
}
Ok(())
}
fn invoke_buffa_build(
&self,
_out_dir: &Path,
effective_includes: &[PathBuf],
) -> Result<(), CodegenError> {
let mut cfg = buffa_build::Config::new()
.files(&self.files)
.includes(effective_includes)
.include_file("obs_buffa.rs")
.generate_views(true);
if let Some(explicit_out) = &self.out_dir {
cfg = cfg.out_dir(explicit_out);
}
if let DescriptorSource::Precompiled(path) = &self.descriptor_source {
cfg = cfg.descriptor_set(path);
}
cfg.compile()
.map_err(|e| CodegenError::Buffa(e.to_string()))?;
Ok(())
}
fn produce_fds(
&self,
out_dir: &Path,
effective_includes: &[PathBuf],
) -> Result<Vec<u8>, CodegenError> {
match &self.descriptor_source {
DescriptorSource::Protoc => self.invoke_protoc(out_dir, effective_includes),
DescriptorSource::Precompiled(path) => {
std::fs::read(path).map_err(CodegenError::DescriptorIo)
}
}
}
fn invoke_protoc(
&self,
out_dir: &Path,
effective_includes: &[PathBuf],
) -> Result<Vec<u8>, CodegenError> {
let protoc = std::env::var("PROTOC").unwrap_or_else(|_| "protoc".to_string());
let fds_path = out_dir.join("obs").join("fds.bin");
let mut cmd = Command::new(&protoc);
cmd.arg("--include_imports");
cmd.arg(format!("--descriptor_set_out={}", fds_path.display()));
for inc in effective_includes {
cmd.arg(format!("--proto_path={}", inc.display()));
}
for f in &self.files {
cmd.arg(f);
}
let status = cmd
.status()
.map_err(|e| CodegenError::Protoc(format!("failed to spawn protoc: {e}")))?;
if !status.success() {
return Err(CodegenError::Protoc(format!("protoc exit status {status}")));
}
std::fs::read(&fds_path).map_err(CodegenError::DescriptorIo)
}
}
fn collect_event_decls(pool: &DescriptorPool) -> Result<Vec<EventDecl>, CodegenError> {
let mut events: Vec<EventDecl> = Vec::new();
for msg in pool.all_messages() {
let dp = msg.descriptor_proto();
if !dp.options.is_set() {
continue;
}
let mut bytes = Vec::new();
dp.options.__buffa_unknown_fields.write_to(&mut bytes);
let Some(event_opts) = read_event_options(&bytes, msg.full_name())? else {
continue;
};
let mut decl = EventDecl {
full_name: msg.full_name().to_string(),
event: event_opts,
fields: Vec::new(),
};
for f in msg.fields() {
let fdp = f.descriptor_proto();
let mut fbytes = Vec::new();
if fdp.options.is_set() {
fdp.options.__buffa_unknown_fields.write_to(&mut fbytes);
}
let opts = read_field_options(&fbytes, &format!("{}/{}", msg.full_name(), f.name()))?
.unwrap_or_default();
let proto_type = Some(map_kind_to_lint_type(&f.kind()));
let wire_rust_type = map_kind_to_rust_type(&f.kind());
let enum_rust_path = match f.kind() {
Kind::Enum(enum_desc) => Some(enum_to_rust_path(&enum_desc)),
_ => None,
};
decl.fields.push(FieldDecl {
name: f.name().to_string(),
number: f.number(),
options: opts,
proto_type,
wire_rust_type,
enum_rust_path,
});
}
events.push(decl);
}
events.sort_by(|a, b| a.full_name.cmp(&b.full_name));
Ok(events)
}
fn enum_to_rust_path(enum_desc: &buffa_reflect::EnumDescriptor) -> String {
let mut parents: Vec<String> = Vec::new();
let mut cursor = enum_desc.parent_message();
while let Some(msg) = cursor {
parents.push(heck::AsSnakeCase(msg.name()).to_string());
cursor = msg.parent_message();
}
parents.reverse();
let file = enum_desc.parent_file();
let package = file.package().trim_start_matches('.');
let mut path = String::new();
if !package.is_empty() {
for seg in package.split('.') {
if !path.is_empty() {
path.push_str("::");
}
path.push_str(seg);
}
}
for p in &parents {
if !path.is_empty() {
path.push_str("::");
}
path.push_str(p);
}
if !path.is_empty() {
path.push_str("::");
}
path.push_str(enum_desc.name());
path
}
fn map_kind_to_rust_type(k: &Kind) -> Option<&'static str> {
match k {
Kind::Bool => Some("bool"),
Kind::Int32 | Kind::Sint32 | Kind::Sfixed32 => Some("i32"),
Kind::Int64 | Kind::Sint64 | Kind::Sfixed64 => Some("i64"),
Kind::Uint32 | Kind::Fixed32 => Some("u32"),
Kind::Uint64 | Kind::Fixed64 => Some("u64"),
Kind::Float => Some("f32"),
Kind::Double => Some("f64"),
_ => None,
}
}
fn map_kind_to_lint_type(k: &Kind) -> LintProtoType {
match k {
Kind::String => LintProtoType::String,
Kind::Bytes => LintProtoType::Bytes,
Kind::Bool => LintProtoType::Bool,
Kind::Double | Kind::Float => LintProtoType::Float,
Kind::Int32
| Kind::Int64
| Kind::Sint32
| Kind::Sint64
| Kind::Sfixed32
| Kind::Sfixed64 => LintProtoType::SignedInteger,
Kind::Uint32 | Kind::Uint64 | Kind::Fixed32 | Kind::Fixed64 => {
LintProtoType::UnsignedInteger
}
Kind::Enum(_) => LintProtoType::Other("enum".to_string()),
Kind::Message(_) => LintProtoType::Other("message".to_string()),
_ => LintProtoType::Other("unknown".to_string()),
}
}
pub const EMBEDDED_OPTIONS_PROTO: &str = include_str!("../proto/obs/v1/options.proto");
pub const EMBEDDED_ENUMS_PROTO: &str = include_str!("../proto/obs/v1/enums.proto");
pub fn materialise_embedded_options(dir: &Path) -> std::io::Result<()> {
let target = dir.join("obs").join("v1");
std::fs::create_dir_all(&target)?;
std::fs::write(target.join("options.proto"), EMBEDDED_OPTIONS_PROTO)?;
std::fs::write(target.join("enums.proto"), EMBEDDED_ENUMS_PROTO)?;
Ok(())
}