substrait_explain/
json.rs

1//! The standard `serde`/pbjson JSON encoding used by Rust stores `google.protobuf.Any`
2//! fields as `{"typeUrl": "...", "value": "<base64>"}`. Go's `protojson` library uses a
3//! different encoding: `{"@type": "...", "field1": val, ...}` where the concrete message's
4//! fields are inlined. `serde_json::from_str::<Plan>` fails on Go-produced JSON because it
5//! only understands the `typeUrl/value` form.
6//!
7//! [`prost_reflect::DynamicMessage`] implements the full protobuf JSON mapping spec and
8//! handles both forms, as long as the `DescriptorPool` contains the schema for every type
9//! URL referenced in the JSON.
10//!
11//! This module exposes [`build_descriptor_pool`] (to construct
12//! the pool, optionally merging in extra descriptor blobs for extension types) and
13//! [`parse_json`] (to parse a JSON string into a [`Plan`] using the pool).
14//!
15//! # Example
16//!
17//! ```rust,ignore
18//! use substrait_explain::json::{build_descriptor_pool, parse_json};
19//!
20//! static MY_EXT: &[u8] = include_bytes!("my_extensions.bin");
21//! let pool = build_descriptor_pool(&[MY_EXT]).unwrap();
22//!  // Works with both Go protojson and Rust pbjson encoding.
23//! let plan = parse_json(json_str, &pool).unwrap();
24//! ```
25
26use std::collections::HashSet;
27
28use anyhow::Context;
29use prost::Message;
30use prost_reflect::{DescriptorPool, DynamicMessage};
31use prost_types::FileDescriptorSet;
32use substrait::proto::{FILE_DESCRIPTOR_SET, Plan};
33
34/// Build a [`DescriptorPool`] covering the Substrait core schema plus any extra
35/// descriptor passed in.
36pub fn build_descriptor_pool(extra_descriptors: &[&[u8]]) -> anyhow::Result<DescriptorPool> {
37    let mut fds = FileDescriptorSet::decode(FILE_DESCRIPTOR_SET)
38        .context("failed to decode substrait core descriptor")?;
39
40    // Descriptor blobs compiled from proto files bundle their transitive dependencies,
41    // therefore custom descriptors are likely to have repeat file names
42    // such as: google/protobuf/timestamp.proto, google/protobuf/any.proto,
43    // which are also present in substrait core protos.
44    // DescriptorPool::decode treats duplicate filenames as a hard error.
45    // Track filenames already in the set so we can skip duplicates.
46    let mut seen: HashSet<String> = fds.file.iter().map(|f| f.name().to_owned()).collect();
47
48    for blob in extra_descriptors {
49        let extra =
50            FileDescriptorSet::decode(*blob).context("failed to decode extra descriptor")?;
51        for f in extra.file {
52            if seen.insert(f.name().to_owned()) {
53                fds.file.push(f);
54            }
55        }
56    }
57
58    DescriptorPool::decode(fds.encode_to_vec().as_slice())
59        .context("failed to build descriptor pool")
60}
61
62/// - **Naive** (`{"typeUrl": "...", "value": "<base64>"}`): decoded via
63///   `serde_json` and `pbjson`.
64///   - This takes the protobuf fields of an `Any` (`type_url`, `value`) and
65///     serializes them like it would any other field. This is the 'naive'
66///     approach to JSON encoding protobufs; see
67///     <https://github.com/influxdata/pbjson/issues/2>
68/// - **Standard** (`{"@type": "...", "field": value, ...}`): decoded via
69///   `prost-reflect`
70///   - `Any` is a Well-Known Type in Protobuf, so in the standard, it has
71///     special handling: the protobuf `type_url` should become the JSON `@type`
72///     field, and other fields should be inlined. See
73///     <https://protobuf.dev/reference/protobuf/google.protobuf/#any>.
74///   - This requires the concrete type's schema to be present in `pool`.
75///
76/// The naive method is tried first (via `serde_json` + `pbjson`); we fall back
77/// to `prost-reflect`, which requires descriptors but can decode
78/// standards-correct JSON-encoded protobufs.
79pub fn parse_json(json: &str, pool: &DescriptorPool) -> anyhow::Result<Plan> {
80    // serde handles the parsing of rust pbjson
81    if let Ok(plan) = serde_json::from_str::<Plan>(json) {
82        return Ok(plan);
83    }
84
85    //  prost-reflect's JSON deserializer handles google.protobuf.Any specifically.
86    //   DynamicMessage::deserialize implements the proto3 JSON mapping
87    //   spec, which encode Any as: { "@type": "type.googleapis.com/pkg.Msg", "field1": val, ... }
88    let plan_desc = pool
89        .get_message_by_name("substrait.Plan")
90        .context("substrait.Plan not found in descriptor pool")?;
91
92    let dyn_msg =
93        DynamicMessage::deserialize(plan_desc, &mut serde_json::Deserializer::from_str(json))
94            .context("failed to parse JSON as substrait.Plan")?;
95
96    Plan::decode(dyn_msg.encode_to_vec().as_slice())
97        .context("failed to decode Plan from dynamic message bytes")
98}
substrait_explain/json.rs

substrait_explain/
json.rs
