olai_codegen/codegen/mod.rs
1//! Code generation module for REST API handlers and language bindings.
2//!
3//! This module provides the core functionality for generating Rust code from
4//! protobuf metadata extracted from service definitions.
5//!
6//! ## Architecture
7//!
8//! The code generation process follows these phases:
9//! 1. **Analysis**: Process collected metadata to understand service structure
10//! 2. **Planning**: Determine what code needs to be generated
11//! 3. **Generation**: Create Rust code using templates and metadata
12//! 4. **Output**: Write generated code to appropriate files
13//!
14//! ## Generated Code Types
15//!
16//! - **Handler Traits**: Async trait definitions for service operations
17//! - **Request Extractors**: Axum FromRequest/FromRequestParts implementations
18//! - **Route Handlers**: Axum handler functions that delegate to traits
19//! - **Client Code**: HTTP client implementations for services
20//! - **Type Mappings**: Conversions between protobuf and Rust types
21
22use std::collections::{BTreeMap, BTreeSet, HashMap};
23
24use convert_case::{Case, Casing};
25use itertools::Itertools;
26use proc_macro2::TokenStream;
27use quote::{format_ident, quote};
28use syn::Ident;
29
30use crate::error::{Error, Result};
31
32use crate::analysis::{
33 BodyField, GenerationPlan, ManagedResource, MethodPlan, RequestParam, RequestType, ServicePlan,
34 analyze_metadata, split_body_fields,
35};
36use crate::output;
37use crate::parsing::types::{self, RenderContext, UnifiedType};
38use crate::parsing::{CodeGenMetadata, MessageField, MessageInfo};
39
40mod aggregate;
41mod bindings;
42mod builder;
43mod client;
44mod client_handler;
45mod config;
46mod connect_client;
47mod error_module;
48mod handler;
49mod json_int64;
50pub(crate) mod node;
51mod python;
52mod resource_client;
53mod resources;
54mod server;
55mod tokens;
56mod transport;
57mod wasm;
58
59pub use config::{
60 BindingsConfig, ClientHandlerImpls, ClientProtocol, ClientProtocols, CodeGenConfig,
61 CodeGenOutput, DEFAULT_TRANSPORT_TYPE_PATH, ErrorModuleConfig, ModelsPath, PyPackageConfig,
62 PythonOptions, RestClientOptions, Runtime,
63};
64pub(crate) use tokens::{doc_tokens, format_tokens, format_tokens_static};
65pub use transport::REQWEST_TRANSPORT_TYPE_PATH;
66
67/// How a language binding lowers a method's call into the underlying Rust client.
68///
69/// Resource services expose a scoped client (e.g. `CatalogClient`) that captures path params;
70/// resource-less services have no such client, so their methods live on the root client and must
71/// pass every param themselves (the "flat" shape, treated like collection methods).
72#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73pub(crate) enum BindingMode {
74 /// Method lives on a resource-scoped client that already holds the path params.
75 Scoped,
76 /// Method lives on the root client and receives all params, including path params.
77 Flat,
78}
79
80/// Language-agnostic description of a service's resource accessor method on the aggregate client.
81///
82/// Computed once from the IR (`ServicePlan::resource_accessor_params`) and rendered by each language
83/// emitter, so the param list and nested-vs-flat decision are shared by construction.
84pub(crate) struct AccessorSpec {
85 /// The resource singular, used as the accessor method name (e.g. `catalog`).
86 pub(crate) singular: String,
87 /// The ordered accessor params (ancestors + `<singular>_name` leaf), e.g.
88 /// `["catalog_name", "schema_name"]` or `["name"]`.
89 pub(crate) params: Vec<String>,
90 /// Whether the resource is nested (has parent components); its full name is the dot-joined
91 /// params and the accessor delegates through a `<singular>_from_full_name` helper.
92 pub(crate) nested: bool,
93}
94
95impl AccessorSpec {
96 /// The `format!` template for joining the params into a composite full name
97 /// (e.g. `"{}.{}"` for two params). Only meaningful when [`Self::nested`] is true.
98 pub(crate) fn join_format(&self) -> String {
99 std::iter::repeat_n("{}", self.params.len())
100 .collect::<Vec<_>>()
101 .join(".")
102 }
103}
104
105impl MethodPlan {
106 pub(crate) fn resource_client_method(&self) -> Ident {
107 match &self.scoped_verb {
108 Some(verb) => format_ident!("{}", verb),
109 None => format_ident!("{}", self.handler_function_name),
110 }
111 }
112
113 pub(crate) fn base_method_ident(&self) -> Ident {
114 format_ident!("{}", self.handler_function_name)
115 }
116}
117
118/// Generate all code described by `config` from `metadata`.
119///
120/// Writes the following outputs, depending on which [`CodeGenOutput`] fields are `Some`:
121///
122/// | Field | Contents |
123/// |-------|----------|
124/// | `output.common` | Axum extractor code, per-service `mod.rs` (always written) |
125/// | `output.models_gen` | `labels.rs` resource-enum file (falls back to `common` if `None`) |
126/// | `output.server` | Handler trait + Axum route wiring per service |
127/// | `output.client` | HTTP client structs and request builders per service |
128/// | `output.python` | PyO3 binding wrappers |
129/// | `output.node` | NAPI binding wrappers |
130/// | `output.node_ts` | TypeScript client (`client.ts`) |
131///
132/// # Required fields
133///
134/// - `output.common` is always required.
135/// - `bindings` must be `Some` when any of `output.python`, `output.node`, or `output.node_ts`
136/// is `Some`; otherwise returns [`Error::MissingBindingsConfig`].
137/// - `models_path_template` and `models_path_crate_template` must be valid Rust path templates
138/// (containing `{service}`); invalid templates return [`Error::InvalidModelsPathTemplate`].
139///
140/// # Optional fields
141///
142/// Setting `generate_resource_enum` to `false` skips `labels.rs` generation.
143/// Setting `bindings` to `None` skips all language-binding output.
144///
145/// Call [`CodeGenConfig::validate`] before this function to surface config errors at
146/// construction time rather than mid-generation.
147///
148/// # Errors
149///
150/// Returns:
151///
152/// - [`Error::InvalidModelsPathTemplate`] if `models_path_template` or `models_path_crate_template`
153/// is not a valid Rust path template.
154/// - [`Error::InvalidConfigPath`] if `context_type_path` or `result_type_path` is not a valid Rust
155/// path.
156/// - [`Error::MissingBindingsConfig`] if any of the `python`, `node`, or `node_ts` outputs is
157/// requested without a `bindings` config.
158/// - [`Error::Build`] if a REST client or Python option lacks an option it needs.
159/// - [`Error::MissingAnnotation`] or [`Error::Build`] surfaced while analyzing the metadata (e.g. a
160/// required `google.api` annotation is missing, or a service cannot be planned).
161/// - [`Error::InvalidRustPath`] if a proto message or field maps to an invalid Rust path, and
162/// [`Error::InvalidErrorTypePath`] if the configured `error_type_path` is not a valid Rust path.
163/// - [`Error::GeneratedParse`] if an internal generator produces tokens that do not parse as a Rust
164/// file.
165/// - [`Error::Io`] if creating the models directory or writing any generated file fails.
166pub fn generate_code(metadata: &CodeGenMetadata, config: &CodeGenConfig) -> Result<()> {
167 // Validate templates early so callers get a clean error before any generation starts.
168 let models_path = ModelsPath::new(&config.models_path_template)?;
169 let models_path_crate = ModelsPath::new(&config.models_path_crate_template)?;
170
171 // The context/result type paths come from user config and are parsed as `syn::Path` deep in
172 // the per-method generators (handler/client/server/builder). Validate them once here so a
173 // typo surfaces as a clean error instead of panicking mid-generation.
174 for (field, path) in [
175 ("context_type", &config.context_type_path),
176 ("result_type", &config.result_type_path),
177 ] {
178 syn::parse_str::<syn::Path>(path).map_err(|source| Error::InvalidConfigPath {
179 field,
180 path: path.clone(),
181 source,
182 })?;
183 }
184
185 // Validate that bindings config is present when language-binding output is requested.
186 if (config.output.python.is_some()
187 || config.output.node.is_some()
188 || config.output.node_ts.is_some())
189 && config.bindings.is_none()
190 {
191 return Err(Error::MissingBindingsConfig);
192 }
193
194 // The Connect client emitter needs the connect-rust client's import path.
195 if config.client_protocols.connect
196 && config.output.client.is_some()
197 && config.connect_client_path.is_none()
198 {
199 return Err(Error::MissingConnectClientPath);
200 }
201
202 config.validate_options()?;
203
204 let plan = analyze_metadata(metadata)?;
205
206 // Validate every service's `{service}` substitution up front so a base path that isn't a
207 // valid Rust path segment fails cleanly here rather than panicking inside per-service codegen.
208 for service in &plan.services {
209 models_path.validate_for(&service.base_path)?;
210 models_path_crate.validate_for(&service.base_path)?;
211 }
212
213 if let Some(ref common_dir) = config.output.common {
214 let common_code = generate_common_code(&plan, metadata, config)?;
215 output::write_generated_code(&common_code, common_dir)?;
216 }
217
218 if config.output.models.is_some() {
219 let subdir = config
220 .output
221 .models_subdir_path()
222 .expect("models is Some so subdir is Some");
223 std::fs::create_dir_all(&subdir).map_err(Error::Io)?;
224
225 if config.generate_resource_enum {
226 let resource_enum = resources::generate_resource_enum(
227 &plan,
228 metadata,
229 config,
230 config.error_type_path.as_deref(),
231 )?;
232 let mut models_files = GeneratedCode {
233 files: std::collections::HashMap::new(),
234 };
235 models_files
236 .files
237 .insert("labels.rs".to_string(), resource_enum);
238 output::write_generated_code(&models_files, &subdir)?;
239 }
240
241 // When Python bindings are generated, the bindings pass model types
242 // across the FFI boundary by their bare type, so the model crate must
243 // carry `FromPyObject`/`IntoPyObject` impls. Emit them here (orphan rule:
244 // the model types are local to this crate). Gated on the `python` feature
245 // in the generated `mod.rs`.
246 let mut include_pyo3_impls = false;
247 if config.output.python.is_some() {
248 let pyo3_impls = python::generate_pyo3_impls(&plan, metadata, config)?;
249 if !pyo3_impls.is_empty() {
250 let mut pyo3_files = GeneratedCode {
251 files: std::collections::HashMap::new(),
252 };
253 pyo3_files
254 .files
255 .insert("pyo3_impls.rs".to_string(), pyo3_impls);
256 output::write_generated_code(&pyo3_files, &subdir)?;
257 include_pyo3_impls = true;
258 }
259 }
260
261 let include_labels = config.generate_resource_enum;
262 // The Axum extractors are co-located in this same directory (the common
263 // generator wrote `<svc>/server.rs` + `<svc>/mod.rs` here), and this
264 // `mod.rs` overwrites the common generator's top-level `mod.rs` — so it
265 // re-declares the per-service extractor submodules it emitted.
266 let mut mod_content = generate_models_mod(
267 &plan.services,
268 include_labels,
269 include_pyo3_impls,
270 config.output.common.is_some(),
271 metadata,
272 config.runtime,
273 );
274 if config.rest_client.json_int64_as_number {
275 mod_content.push_str(&json_int64::module());
276 if metadata_has_json_fields(metadata) {
277 mod_content.push_str(&json_int64::value_module());
278 }
279 }
280 // Route through write_generated_code so the `// @generated` header is prepended like
281 // every other emitted file, rather than writing it raw.
282 let mut mod_files = GeneratedCode {
283 files: std::collections::HashMap::new(),
284 };
285 mod_files.files.insert("mod.rs".to_string(), mod_content);
286 output::write_generated_code(&mod_files, &subdir)?;
287 }
288
289 if let Some(ref server_dir) = config.output.server {
290 let server_code = generate_server_code(&plan, metadata, config)?;
291 output::write_generated_code(&server_code, server_dir)?;
292 }
293
294 if let Some(ref client_dir) = config.output.client {
295 let client_code = generate_client_code(&plan, metadata, config)?;
296 output::write_generated_code(&client_code, client_dir)?;
297 }
298
299 if let Some(ref python_dir) = config.output.python {
300 let python_code = generate_python_code(&plan, metadata, config)?;
301 output::write_generated_code(&python_code, python_dir)?;
302
303 let package_files = python::package_files(config, &plan, metadata);
304 if !package_files.is_empty() {
305 let files = GeneratedCode {
306 files: package_files.into_iter().collect(),
307 };
308 output::write_generated_code(&files, &config.python_options().package.dir)?;
309 }
310 }
311
312 if let Some(ref node_dir) = config.output.node {
313 let node_code = generate_node_code(&plan, metadata, config)?;
314 output::write_generated_code(&node_code, node_dir)?;
315 }
316
317 if let Some(ref node_ts_dir) = config.output.node_ts {
318 let node_ts_code = generate_node_ts_code(&plan, metadata, config)?;
319 output::write_generated_code(&node_ts_code, node_ts_dir)?;
320 }
321
322 if let Some(ref wasm_dir) = config.output.wasm {
323 let wasm_code = generate_wasm_code(&plan, metadata, config)?;
324 output::write_generated_code(&wasm_code, wasm_dir)?;
325 }
326
327 Ok(())
328}
329
330/// Whether a project message has a free-form JSON field (`google.protobuf.Value`, `Struct`,
331/// `ListValue`), also as a map value or oneof variant.
332pub(crate) fn metadata_has_json_fields(metadata: &CodeGenMetadata) -> bool {
333 metadata
334 .messages
335 .iter()
336 .filter(|(name, _)| !name.starts_with(".google."))
337 .flat_map(|(_, info)| &info.fields)
338 .any(|field| {
339 types::JsonType::of_field(&field.unified_type).is_some()
340 || field
341 .oneof_variants
342 .iter()
343 .flatten()
344 .any(|v| types::JsonType::of_field(&v.field_type).is_some())
345 })
346}
347
348fn generate_common_code(
349 plan: &GenerationPlan,
350 metadata: &CodeGenMetadata,
351 config: &CodeGenConfig,
352) -> Result<GeneratedCode> {
353 let mut files = HashMap::new();
354
355 for service in &plan.services {
356 let handler = ServiceHandler {
357 plan: service,
358 metadata,
359 config,
360 };
361 let server_code = server::generate_common(&handler)?;
362 // Nest the extractors under an `extractors/` subdir so their per-service
363 // module names can't collide with the co-located model package modules
364 // (e.g. a `unitycatalog.agent_skills.v0alpha1` package yields a model
365 // module `agent_skills` that would otherwise clash with the `agent_skills`
366 // service extractor module).
367 files.insert(
368 format!("{EXTRACTORS_SUBDIR}/{}/server.rs", service.base_path),
369 server_code,
370 );
371 let module_code = generate_common_module();
372 files.insert(
373 format!("{EXTRACTORS_SUBDIR}/{}/mod.rs", service.base_path),
374 module_code,
375 );
376 }
377
378 // The `extractors/mod.rs` that declares the per-service extractor submodules.
379 files.insert(
380 format!("{EXTRACTORS_SUBDIR}/mod.rs"),
381 generate_extractors_module(&plan.services),
382 );
383
384 let module_code = main_module(&plan.services);
385 files.insert("mod.rs".to_string(), module_code);
386
387 Ok(GeneratedCode { files })
388}
389
390fn generate_server_code(
391 plan: &GenerationPlan,
392 metadata: &CodeGenMetadata,
393 config: &CodeGenConfig,
394) -> Result<GeneratedCode> {
395 let mut files = HashMap::new();
396
397 for service in &plan.services {
398 let handler = ServiceHandler {
399 plan: service,
400 metadata,
401 config,
402 };
403 let trait_code = handler::generate(&handler)?;
404 files.insert(format!("{}/handler.rs", service.base_path), trait_code);
405 let server_code = server::generate_server(&handler)?;
406 files.insert(format!("{}/server.rs", service.base_path), server_code);
407 let module_code = generate_server_module(service);
408 files.insert(format!("{}/mod.rs", service.base_path), module_code);
409 }
410
411 let module_code = main_module(&plan.services);
412 files.insert("mod.rs".to_string(), module_code);
413
414 Ok(GeneratedCode { files })
415}
416
417fn generate_python_code(
418 plan: &GenerationPlan,
419 metadata: &CodeGenMetadata,
420 config: &CodeGenConfig,
421) -> Result<GeneratedCode> {
422 let mut files = HashMap::new();
423
424 let handlers = plan
425 .services
426 .iter()
427 .map(|service| ServiceHandler {
428 plan: service,
429 metadata,
430 config,
431 })
432 .collect_vec();
433
434 // Per-service scoped client modules are only meaningful for services with a scoped binding.
435 // The methods of resource-less and flat-lowered services live on the root client (emitted by
436 // `main_module`), so they don't get a scoped module of their own.
437 for service in handlers.iter().filter(|s| s.emits_scoped_binding()) {
438 let python_code = python::generate(service)?;
439 files.insert(format!("{}.rs", service.plan.base_path), python_code);
440 }
441
442 let module_code = python::main_module(&handlers)?;
443 files.insert("mod.rs".to_string(), module_code);
444 files.extend(python::helper_files(config));
445
446 Ok(GeneratedCode { files })
447}
448
449fn generate_node_code(
450 plan: &GenerationPlan,
451 metadata: &CodeGenMetadata,
452 config: &CodeGenConfig,
453) -> Result<GeneratedCode> {
454 let mut files = HashMap::new();
455
456 let handlers = plan
457 .services
458 .iter()
459 .map(|service| ServiceHandler {
460 plan: service,
461 metadata,
462 config,
463 })
464 .collect_vec();
465
466 // See `generate_python_code`: only services with a scoped binding get a per-service scoped
467 // module; resource-less and flat-lowered services' methods live on the root client.
468 for service in handlers.iter().filter(|s| s.emits_scoped_binding()) {
469 let napi_code = node::generate(service)?;
470 files.insert(format!("{}.rs", service.plan.base_path), napi_code);
471 }
472
473 let module_code = node::main_module(&handlers)?;
474 files.insert("mod.rs".to_string(), module_code);
475
476 Ok(GeneratedCode { files })
477}
478
479fn generate_node_ts_code(
480 plan: &GenerationPlan,
481 metadata: &CodeGenMetadata,
482 config: &CodeGenConfig,
483) -> Result<GeneratedCode> {
484 let handlers = plan
485 .services
486 .iter()
487 .map(|service| ServiceHandler {
488 plan: service,
489 metadata,
490 config,
491 })
492 .collect_vec();
493
494 let ts_code = node::typescript::generate_client_ts(&handlers);
495 let mut files = HashMap::new();
496 files.insert("client.ts".to_string(), ts_code);
497
498 // The `models/index.ts` barrel re-exports every service's generated protobuf-es modules so
499 // `client.ts`'s `from "./models"` imports (message and service request/response types) resolve.
500 let models_barrel = node::typescript::generate_models_barrel(&handlers);
501 files.insert("models/index.ts".to_string(), models_barrel);
502
503 Ok(GeneratedCode { files })
504}
505
506/// Generate the WASM/browser `#[wasm_bindgen]` binding layer plus its `.d.ts`.
507///
508/// Emits `bindings.rs` (a `#[wasm_bindgen]` wrapper per service client, exchanging request/response
509/// values as plain JS objects via `serde-wasm-bindgen`) and `client.d.ts` (TypeScript declarations
510/// for JS/TS consumers). The wrapped clients are the generated low-level clients, which must be
511/// generated with the WASM transport (`transport_type_path = "olai_http_wasm::WasmClient"`).
512fn generate_wasm_code(
513 plan: &GenerationPlan,
514 metadata: &CodeGenMetadata,
515 config: &CodeGenConfig,
516) -> Result<GeneratedCode> {
517 let handlers = plan
518 .services
519 .iter()
520 .map(|service| ServiceHandler {
521 plan: service,
522 metadata,
523 config,
524 })
525 .collect_vec();
526
527 let mut files = HashMap::new();
528 // Emit as `mod.rs` so the client crate mounts the bindings with a plain
529 // `mod wasm;` (no `#[path]` alias) over the `wasm/` convention subdir.
530 files.insert("mod.rs".to_string(), wasm::generate_bindings(&handlers)?);
531 files.insert("client.d.ts".to_string(), wasm::generate_dts(&handlers));
532
533 Ok(GeneratedCode { files })
534}
535
536fn generate_client_code(
537 plan: &GenerationPlan,
538 metadata: &CodeGenMetadata,
539 config: &CodeGenConfig,
540) -> Result<GeneratedCode> {
541 let mut files = HashMap::new();
542 let protocols = config.client_protocols.selected();
543 let multi = config.client_protocols.is_multi();
544 // Module path and type of each resource-scoped client, for the root re-exports.
545 let mut resource_clients = Vec::new();
546
547 for service in &plan.services {
548 let handler = ServiceHandler {
549 plan: service,
550 metadata,
551 config,
552 };
553
554 // Emit each selected protocol's low-level client + builders. With a single protocol the
555 // files are flat under `<service>/`; with both they go under `<service>/<protocol>/` to
556 // avoid the two `client.rs`/`builders.rs` colliding. The shared builder layer binds to the
557 // protocol's low-level client and only calls `client.<method>(&request)`.
558 let mut emitted_resource_client = false;
559 for &protocol in &protocols {
560 let prefix = if multi {
561 format!("{}/{}", service.base_path, protocol.subdir())
562 } else {
563 service.base_path.clone()
564 };
565
566 let client_code = match protocol {
567 ClientProtocol::Rest => client::generate(&handler)?,
568 ClientProtocol::Connect => connect_client::generate(&handler)?,
569 };
570 files.insert(format!("{prefix}/client.rs"), client_code);
571 files.insert(
572 format!("{prefix}/builders.rs"),
573 builder::generate(&handler, protocol)?,
574 );
575
576 // Ergonomic resource-scoped clients are a REST-only feature (driven by
577 // `google.api.resource` hierarchies). ConnectRPC-only protos rarely carry those, so the
578 // Connect client emits just the flat per-service client + builders.
579 let has_resource_client = protocol == ClientProtocol::Rest
580 && config.output.generate_resource_clients
581 && resource_client::generate(&handler, plan)?
582 .map(|code| {
583 files.insert(format!("{prefix}/resource.rs"), code);
584 })
585 .is_some();
586
587 if has_resource_client && let Some(scoped) = handler.scoped_client_type() {
588 resource_clients.push((prefix.clone(), scoped, handler.client_feature_cfg()));
589 }
590
591 // Handler-trait impls forward to the REST client, so only it gets them.
592 let handler_impl = config
593 .rest_client
594 .handler_impls
595 .as_ref()
596 .filter(|impls| {
597 protocol == ClientProtocol::Rest && !impls.skip.contains(&service.base_path)
598 })
599 .map(|impls| {
600 client_handler::generate(&handler, impls).map(|code| {
601 files.insert(format!("{prefix}/handler.rs"), code);
602 impls.feature.as_str()
603 })
604 })
605 .transpose()?;
606
607 files.insert(
608 format!("{prefix}/mod.rs"),
609 generate_client_module(has_resource_client, handler_impl),
610 );
611 emitted_resource_client |= has_resource_client;
612 }
613
614 // When both protocols are emitted under subdirs, the `<service>/mod.rs` mounts and re-exports
615 // them. (Single-protocol output already wrote a flat `<service>/mod.rs` above.)
616 if multi {
617 files.insert(
618 format!("{}/mod.rs", service.base_path),
619 generate_protocol_split_module(&protocols),
620 );
621 }
622 let _ = emitted_resource_client;
623 }
624
625 // Top-level aggregate root client (e.g. `UnityCatalogClient`), emitted only when a bindings
626 // config (which carries the aggregate name) is present.
627 let has_aggregate = aggregate::generate(plan, metadata, config)?
628 .map(|aggregate_code| {
629 files.insert("client.rs".to_string(), aggregate_code);
630 })
631 .is_some();
632
633 if config.uses_reqwest_transport() {
634 files.insert(
635 "transport.rs".to_string(),
636 transport::generate(config.rest_client.user_agent.as_deref()),
637 );
638 }
639 if let Some(errors) = &config.rest_client.error_module {
640 files.insert(
641 "error.rs".to_string(),
642 error_module::generate(config, errors, metadata)?,
643 );
644 }
645
646 files.insert(
647 "mod.rs".to_string(),
648 generate_client_main_module(plan, metadata, config, has_aggregate, &resource_clients),
649 );
650
651 Ok(GeneratedCode { files })
652}
653
654fn generate_common_module() -> String {
655 let tokens = quote! {
656 #[cfg(feature = "axum")]
657 pub mod server;
658 };
659 format_tokens_static(tokens)
660}
661
662/// Subdirectory (under the co-located models `_gen/` dir) that holds the Axum
663/// request extractors, one `<service>/server.rs` per service. Nesting them keeps
664/// their module names from colliding with the model package modules that share
665/// the dir.
666pub(crate) const EXTRACTORS_SUBDIR: &str = "extractors";
667
668/// The `extractors/mod.rs` that declares one `pub mod <service>;` per service.
669/// Axum-coupled, so gated on the `axum` feature.
670fn generate_extractors_module(services: &[ServicePlan]) -> String {
671 let decls: Vec<TokenStream> = services
672 .iter()
673 .map(|s| format_ident!("{}", s.base_path))
674 // Dedupe: multiple services can share a base_path.
675 .collect::<BTreeSet<_>>()
676 .into_iter()
677 .map(|m| quote! { #[cfg(feature = "axum")] pub mod #m; })
678 .collect();
679 format_tokens_static(quote! { #(#decls)* })
680}
681
682fn generate_server_module(service: &ServicePlan) -> String {
683 let handler_ident = format_ident!("{}", service.handler_name);
684 let tokens = quote! {
685 pub use handler::#handler_ident;
686
687 mod handler;
688 #[cfg(feature = "axum")]
689 pub mod server;
690 };
691 format_tokens_static(tokens)
692}
693
694/// `handler_impl` is the cargo feature gating the service's handler-trait impl module, when one
695/// was generated for it.
696fn generate_client_module(has_resource_client: bool, handler_impl: Option<&str>) -> String {
697 // Emit the scoped resource client module only when it was generated for this service.
698 // Generated client modules structurally trip a few lints: the full client
699 // surface is emitted whether or not a given consumer calls every builder
700 // (`dead_code`), flat methods can exceed the arg-count threshold
701 // (`too_many_arguments`), and proto comment prose may wrap list items
702 // (`doc_lazy_continuation`). Allow these on the module so consumers' own
703 // `-D warnings` builds stay clean; genuine issues are fixed in the generator.
704 let allow = quote! {
705 #[allow(dead_code, clippy::too_many_arguments, clippy::doc_lazy_continuation)]
706 };
707 // Glob re-exports of the per-service submodules. Not every submodule
708 // contributes a name that the parent's consumers use directly, so a
709 // `pub use <mod>::*` can trip `unused_imports` under `-D warnings` (notably
710 // under buffa, where some helper types are unreferenced); allow it.
711 let reexport_allow = quote! { #[allow(unused_imports)] };
712 let resource_module = if has_resource_client {
713 quote! {
714 #reexport_allow
715 pub use resource::*;
716 #allow
717 pub mod resource;
718 }
719 } else {
720 quote! {}
721 };
722 // The impls are trait impls only (nothing to re-export). Client futures aren't `Send` on
723 // wasm32, which `async_trait` handler traits require, so the module is native-only.
724 let handler_module = match handler_impl {
725 Some(feature) => quote! {
726 #[cfg(all(feature = #feature, not(target_arch = "wasm32")))]
727 mod handler;
728 },
729 None => quote! {},
730 };
731 let tokens = quote! {
732 #reexport_allow
733 pub use client::*;
734 #reexport_allow
735 pub use builders::*;
736 #resource_module
737 #handler_module
738
739 #allow
740 pub mod client;
741 #allow
742 pub mod builders;
743 };
744 format_tokens_static(tokens)
745}
746
747/// The `<service>/mod.rs` for a service that emits **both** protocols: mount each protocol's client
748/// as its own submodule (`pub mod rest;` / `pub mod connect;`).
749///
750/// They are *not* glob-re-exported into the service module: both protocols define identically-named
751/// types (`GetThingBuilder`, …), so a flat re-export would collide. Callers disambiguate by module
752/// (`read::rest::GetThingBuilder` vs `read::connect::GetThingBuilder`).
753fn generate_protocol_split_module(protocols: &[ClientProtocol]) -> String {
754 let allow = quote! {
755 #[allow(dead_code, clippy::too_many_arguments, clippy::doc_lazy_continuation)]
756 };
757 let mods: Vec<TokenStream> = protocols
758 .iter()
759 .map(|p| {
760 let m = format_ident!("{}", p.subdir());
761 quote! { #allow pub mod #m; }
762 })
763 .collect();
764 format_tokens_static(quote! { #(#mods)* })
765}
766
767pub fn main_module(services: &[ServicePlan]) -> String {
768 let service_modules: Vec<TokenStream> = services
769 .iter()
770 .map(|s| {
771 let module_name = format_ident!("{}", s.base_path);
772 quote! { pub mod #module_name; }
773 })
774 .collect();
775
776 let tokens = quote! {
777 #(#service_modules)*
778 };
779 format_tokens_static(tokens)
780}
781
782/// The client root `mod.rs`. It re-exports the `resource_clients` (module path, type and the
783/// service's feature gate) so a crate can export the whole client with `pub use codegen::*;`.
784fn generate_client_main_module(
785 plan: &GenerationPlan,
786 metadata: &CodeGenMetadata,
787 config: &CodeGenConfig,
788 has_aggregate: bool,
789 resource_clients: &[(String, Ident, TokenStream)],
790) -> String {
791 let services = &plan.services;
792 let result_path = config.client_result_path();
793 let service_modules: Vec<TokenStream> = services
794 .iter()
795 .map(|s| {
796 let module_name = format_ident!("{}", s.base_path);
797 let cfg = ServiceHandler {
798 plan: s,
799 metadata,
800 config,
801 }
802 .client_feature_cfg();
803 quote! { #cfg pub mod #module_name; }
804 })
805 .collect();
806 let resource_exports: Vec<TokenStream> = resource_clients
807 .iter()
808 .map(|(path, scoped, cfg)| {
809 let segments = path.split('/').map(|segment| format_ident!("{}", segment));
810 quote! { #cfg pub use #(#segments::)*#scoped; }
811 })
812 .collect();
813
814 // Export the aggregate root client (`client.rs`) when it was generated. Per-service files live
815 // under `<base_path>/`, so a top-level `client` module never collides with them.
816 // The aggregate client has many flat methods (arg-count) and renders proto
817 // comment prose (doc list wrapping), so allow those structural lints.
818 let aggregate_module = if has_aggregate {
819 quote! {
820 #[allow(clippy::too_many_arguments, clippy::doc_lazy_continuation)]
821 pub mod client;
822 #[allow(unused_imports)]
823 pub use client::*;
824 }
825 } else {
826 quote! {}
827 };
828
829 // The `stream_paginated` helper (and its `futures::Future` import) is used
830 // only by `List` builders' `into_stream`. Emit it solely when some service
831 // has a List method, so it isn't dead code otherwise.
832 let has_list = services
833 .iter()
834 .flat_map(|s| s.methods.iter())
835 .any(|m| matches!(m.request_type, RequestType::List));
836 // With per-resource features, every List-bearing service may be disabled.
837 let pagination_allow = if config.rest_client.resource_features {
838 quote! { #[allow(dead_code)] }
839 } else {
840 quote! {}
841 };
842 let pagination = if has_list {
843 quote! {
844 use futures::Future;
845
846 #pagination_allow
847 pub(super) fn stream_paginated<F, Fut, S, T>(
848 state: S,
849 op: F,
850 ) -> impl futures::Stream<Item = #result_path<T>>
851 where
852 F: Fn(S, Option<String>) -> Fut + Copy,
853 Fut: Future<Output = #result_path<(T, S, Option<String>)>>,
854 {
855 enum PaginationState<T> {
856 Start(T),
857 HasMore(T, String),
858 Done,
859 }
860
861 futures::stream::unfold(
862 PaginationState::Start(state),
863 move |state| async move {
864 let (s, page_token) = match state {
865 PaginationState::Start(s) => (s, None),
866 PaginationState::HasMore(s, page_token) if !page_token.is_empty() => {
867 (s, Some(page_token))
868 }
869 _ => {
870 return None;
871 }
872 };
873
874 let (resp, s, continuation) = match op(s, page_token).await {
875 Ok(resp) => resp,
876 Err(e) => return Some((Err(e), PaginationState::Done)),
877 };
878
879 let next_state = match continuation {
880 Some(token) => PaginationState::HasMore(s, token),
881 None => PaginationState::Done,
882 };
883
884 Some((Ok(resp), next_state))
885 },
886 )
887 }
888 }
889 } else {
890 quote! {}
891 };
892
893 let transport_module = config
894 .uses_reqwest_transport()
895 .then(|| quote! { pub mod transport; });
896 let error_module = config
897 .rest_client
898 .error_module
899 .is_some()
900 .then(|| quote! { pub mod error; });
901 let encode_path_value = services
902 .iter()
903 .flat_map(|s| &s.methods)
904 .any(MethodPlan::formats_path_values)
905 .then(|| {
906 let helper = client::encode_path_value_fn();
907 // With per-resource features, every service that formats path values may be off.
908 quote! { #pagination_allow #helper }
909 });
910
911 let tokens = quote! {
912 #(#service_modules)*
913
914 #(#resource_exports)*
915
916 #aggregate_module
917 #transport_module
918 #error_module
919
920 #pagination
921 #encode_path_value
922 };
923 format_tokens_static(tokens)
924}
925
926/// Generate the `mod.rs` for `crates/common/src/models/`.
927///
928/// Emits `pub mod <service> { pub mod <version> { include!("./<pkg>.rs") } }` blocks
929/// for every service in the plan, plus static re-exports and module declarations.
930/// The model files are co-located in this directory (written by the buf plugin), so
931/// the includes are plain sibling references.
932///
933/// When `include_labels` is `true`, also emits `pub mod labels; pub use labels::{ObjectLabel, Resource};`.
934/// When `include_extractors` is `true`, also declares `pub mod extractors;`.
935///
936/// `metadata` is used to discover all resource-annotated messages (even those not returned
937/// directly by an RPC) so they can be included in `pub use` re-exports.
938pub fn generate_models_mod(
939 services: &[ServicePlan],
940 include_labels: bool,
941 include_pyo3_impls: bool,
942 include_extractors: bool,
943 metadata: &CodeGenMetadata,
944 runtime: Runtime,
945) -> String {
946 // Services are keyed by proto package, not by service: multiple services can share a
947 // package (e.g. `CatalogService` and `SchemaService` both in `example.catalog.v1`), and
948 // emitting one `pub mod`/`pub use` block per service would produce duplicate modules and
949 // imports (E0428). Group by package first, then emit once per package.
950 //
951 // package = "unitycatalog.catalogs.v1"
952 // parts = ["unitycatalog", "catalogs", "v1"]
953 // service module = second-to-last segment ("catalogs"); version = last segment ("v1").
954 let mod_segments = |svc: &ServicePlan| -> (String, String) {
955 let parts: Vec<&str> = svc.package.split('.').collect();
956 if parts.len() >= 2 {
957 (
958 parts[parts.len() - 2].to_string(),
959 parts[parts.len() - 1].to_string(),
960 )
961 } else {
962 (svc.base_path.clone(), "v1".to_string())
963 }
964 };
965
966 // One entry per package, in sorted order, with the union of re-exported type names.
967 let mut packages: BTreeMap<String, (String, String, BTreeSet<String>)> = BTreeMap::new();
968 for svc in services {
969 let (svc_seg, ver_seg) = mod_segments(svc);
970 let entry = packages
971 .entry(svc.package.clone())
972 .or_insert_with(|| (svc_seg, ver_seg, BTreeSet::new()));
973 let type_names = &mut entry.2;
974
975 // Re-export managed resources for this service...
976 type_names.extend(svc.managed_resources.iter().map(|r| r.type_name.clone()));
977
978 // ...plus every resource-annotated message in the same package (catches nested types
979 // like `Column` that aren't direct RPC return types but still carry annotations).
980 let fq_prefix = format!(".{}.", svc.package);
981 let fq_pkg = format!(".{}", svc.package);
982 for (fq_name, msg_info) in &metadata.messages {
983 if msg_info.resource_descriptor.is_some()
984 && (fq_name.starts_with(&fq_prefix) || fq_name.starts_with(&fq_pkg))
985 {
986 let simple = fq_name
987 .rfind('.')
988 .map(|i| &fq_name[i + 1..])
989 .unwrap_or(fq_name.as_str());
990 type_names.insert(simple.to_string());
991 }
992 }
993 }
994
995 let mut service_mods: Vec<TokenStream> = Vec::new();
996 let mut reexports: Vec<TokenStream> = Vec::new();
997 for (package, (svc_seg, ver_seg, type_names)) in &packages {
998 let svc_mod = format_ident!("{}", svc_seg);
999 let ver_mod = format_ident!("{}", ver_seg);
1000
1001 // Models are co-located: the buf plugin writes `<package>.rs` (and, for
1002 // prost, `<package>.serde.rs` / `<package>.tonic.rs`) into this same
1003 // directory, so the include is a plain sibling reference.
1004 let main_include = format!("./{package}.rs");
1005 // The separate `<package>.tonic.rs` (grpc service stubs over the prost
1006 // `ProstCodec`) is only emitted by the prost model plugin. The buffa plugin
1007 // produces no such file, so including it would break the build whenever the
1008 // `grpc` feature is enabled (e.g. CI's `--all-features`). Gate the include on
1009 // the runtime, not just the `grpc` cfg.
1010 let tonic_include_tokens = match runtime {
1011 Runtime::Prost => {
1012 let tonic_include = format!("./{package}.tonic.rs");
1013 quote! {
1014 #[cfg(feature = "grpc")]
1015 include!(#tonic_include);
1016 }
1017 }
1018 Runtime::Buffa => quote! {},
1019 };
1020 service_mods.push(quote! {
1021 pub mod #svc_mod {
1022 pub mod #ver_mod {
1023 include!(#main_include);
1024 #tonic_include_tokens
1025 }
1026 }
1027 });
1028
1029 for type_name in type_names {
1030 let type_ident = format_ident!("{}", type_name);
1031 reexports.push(quote! {
1032 pub use #svc_mod::#ver_mod::#type_ident;
1033 });
1034 }
1035 }
1036
1037 let labels_decl: TokenStream = if include_labels {
1038 quote! {
1039 pub mod labels;
1040 pub use labels::{ObjectLabel, Resource};
1041 }
1042 } else {
1043 quote! {}
1044 };
1045
1046 // PyO3 boundary conversions for the model types (see `python::generate_pyo3_impls`).
1047 // Gated on the `python` feature so non-Python consumers never pull in pyo3.
1048 let pyo3_impls_decl: TokenStream = if include_pyo3_impls {
1049 quote! {
1050 // pyo3's `#[pyclass(from_py_object)]` extracts via `Clone::clone`, which trips
1051 // `clippy::clone_on_copy` on the (`Copy`) enum wrappers. The call lives in the
1052 // macro-generated impl, so only a module-level allow can reach it.
1053 #[cfg(feature = "python")]
1054 #[allow(clippy::clone_on_copy)]
1055 pub mod pyo3_impls;
1056 // Re-export the `Py*` wrapper `#[pyclass]` types at the models root so the
1057 // generated language bindings (and consumers) can name them without
1058 // reaching into the `pyo3_impls` submodule.
1059 #[cfg(feature = "python")]
1060 pub use pyo3_impls::*;
1061 }
1062 } else {
1063 quote! {}
1064 };
1065
1066 // The Axum extractors, when generated, are co-located in this directory under an
1067 // `extractors/` subdir (the common generator wrote `extractors/<svc>/...`
1068 // here), and this `mod.rs` overwrites the top-level `mod.rs` the common
1069 // generator emitted — so it must re-declare the `extractors` module.
1070 // Axum-coupled, so gate it on the `axum` feature. Without this, the
1071 // `FromRequest`/`FromRequestParts` impls never compile and the generated
1072 // route fns can't be mounted. Nesting under `extractors` also avoids the
1073 // per-service module names clashing with the model package modules above.
1074 let common_mods: TokenStream = if include_extractors {
1075 let extractors_mod = format_ident!("{EXTRACTORS_SUBDIR}");
1076 quote! {
1077 #[cfg(feature = "axum")]
1078 pub mod #extractors_mod;
1079 }
1080 } else {
1081 quote! {}
1082 };
1083
1084 // Lints to silence across this module and every `include!`d model file below.
1085 // These are inner attributes, so they cover the descendant `pub mod <svc>::<v1>`
1086 // modules the per-package files are pasted into.
1087 let lint_allows = match runtime {
1088 // The buffa model files carry no allows of their own, so silence their lints here,
1089 // as buffa's own include file does.
1090 Runtime::Buffa => quote! {
1091 #![allow(non_camel_case_types)]
1092 #![allow(clippy::derivable_impls)]
1093 #![allow(clippy::match_single_binding)]
1094 #![allow(clippy::uninlined_format_args)]
1095 #![allow(clippy::doc_lazy_continuation)]
1096 },
1097 Runtime::Prost => quote! {},
1098 };
1099
1100 let tokens = quote! {
1101 // The per-package modules below gate a `tonic.rs` include on
1102 // `#[cfg(feature = "grpc")]`. Consumers that don't define a `grpc`
1103 // feature would otherwise see `unexpected_cfgs`; allow it here.
1104 #![allow(unexpected_cfgs)]
1105 // The grpc-gated `tonic.rs` includes carry empty (`///`) doc comments on
1106 // generated service methods; allow them so `clippy -D warnings` stays
1107 // green when the `grpc` feature is enabled.
1108 #![allow(clippy::empty_docs)]
1109 // The prost-serde plugin emits `write!(f, "...", &FIELDS)` for the
1110 // deserializer's expected-fields message; newer clippy flags the `&` as a
1111 // useless borrow in a formatting macro. Allow it here so plugin output
1112 // stays `-D warnings` clean without post-processing each generated serde file.
1113 #![allow(clippy::useless_borrows_in_formatting)]
1114 #lint_allows
1115
1116 use std::collections::HashMap;
1117
1118 #labels_decl
1119
1120 #pyo3_impls_decl
1121
1122 #(#reexports)*
1123
1124 pub type PropertyMap = HashMap<String, serde_json::Value>;
1125
1126 #(#service_mods)*
1127
1128 #common_mods
1129 };
1130
1131 format_tokens_static(tokens)
1132}
1133
1134/// An in-memory bundle of generated source files awaiting output.
1135///
1136/// Each generator phase (common, server, client, bindings, …) returns one of these; the
1137/// [`crate::output`] writer then materializes the [`files`](Self::files) onto disk, prepending the
1138/// `// @generated` header to each. Keeping generation and writing separate lets generation stay
1139/// pure (no I/O) and makes the produced files easy to assert on in tests.
1140#[derive(Debug)]
1141pub struct GeneratedCode {
1142 /// Generated files keyed by path relative to the phase's output directory.
1143 pub files: HashMap<String, String>,
1144}
1145
1146impl CodeGenMetadata {
1147 fn get_message_meta(&self, message_name: &str) -> Option<MessageMeta<'_>> {
1148 self.messages.get(message_name).map(|info| MessageMeta {
1149 info,
1150 metadata: self,
1151 })
1152 }
1153}
1154
1155/// A protobuf message paired with the metadata needed to resolve its references.
1156///
1157/// Wraps a single [`MessageInfo`](crate::parsing::MessageInfo) together with a borrow of the whole
1158/// [`CodeGenMetadata`] it came from, so the generators can look up a message's fields while still
1159/// being able to resolve types that point at other messages in the same descriptor set. Obtained
1160/// via `CodeGenMetadata::get_message_meta`.
1161pub(crate) struct MessageMeta<'a> {
1162 info: &'a MessageInfo,
1163 #[allow(dead_code)]
1164 metadata: &'a CodeGenMetadata,
1165}
1166
1167/// The per-service generation context shared by every code generator.
1168///
1169/// Bundles the three borrows each generator phase needs for one service: its analyzed
1170/// [`ServicePlan`](crate::analysis::ServicePlan), the full [`CodeGenMetadata`] (for resolving
1171/// types referenced from other services), and the active [`CodeGenConfig`]. The various
1172/// `codegen::*` submodules take a `&ServiceHandler` and read these plus the derived accessors on
1173/// this type (client type names, models paths, [`Self::methods`]) rather than threading the three
1174/// borrows separately. Methods are exposed as [`MethodHandler`]s via [`Self::methods`].
1175pub(crate) struct ServiceHandler<'a> {
1176 pub(crate) plan: &'a ServicePlan,
1177 pub(crate) metadata: &'a CodeGenMetadata,
1178 pub(crate) config: &'a CodeGenConfig,
1179}
1180
1181impl ServiceHandler<'_> {
1182 pub(crate) fn resource(&self) -> Option<&ManagedResource> {
1183 self.plan.managed_resources.first()
1184 }
1185
1186 /// Whether this service manages a `google.api.resource` and therefore gets a resource-scoped
1187 /// client (e.g. `CatalogClient` bound to a name).
1188 ///
1189 /// Resource-less services (entity tag assignments, temporary credentials, delta commits) have
1190 /// no scoped client; their methods are emitted on the root client and pass every param
1191 /// directly ([`BindingMode::Flat`]). See [`Self::binding_mode`].
1192 pub(crate) fn is_resource_scoped(&self) -> bool {
1193 self.resource().is_some()
1194 }
1195
1196 /// The [`BindingMode`] used to lower this service's methods in language bindings.
1197 ///
1198 /// A resource service whose identity can't be modeled by a string-component scoped client
1199 /// (see [`Self::supports_scoped_client`]) is lowered flat, exactly like a resource-less
1200 /// service: every method lives on the root client and passes all params directly.
1201 pub(crate) fn binding_mode(&self) -> BindingMode {
1202 if self.emits_scoped_binding() {
1203 BindingMode::Scoped
1204 } else {
1205 BindingMode::Flat
1206 }
1207 }
1208
1209 /// Whether language bindings emit a per-service scoped client module + accessor for this
1210 /// service. True only for resource services whose identity a string-component scoped client
1211 /// can express; others are lowered flat onto the root client. See [`Self::binding_mode`].
1212 pub(crate) fn emits_scoped_binding(&self) -> bool {
1213 self.is_resource_scoped() && self.supports_scoped_client()
1214 }
1215
1216 /// The language-agnostic spec for this service's resource accessor on the aggregate client, or
1217 /// `None` for resource-less services.
1218 ///
1219 /// All four accessor emitters (Rust aggregate, Python, Node, TypeScript) render from this single
1220 /// spec rather than re-deriving the param list and nesting decision, which previously diverged
1221 /// between the Rust aggregate and the bindings.
1222 pub(crate) fn accessor_spec(&self) -> Option<AccessorSpec> {
1223 let resource = self.resource()?;
1224 let params = self.plan.resource_accessor_params.clone()?;
1225 // A resource is nested when its accessor params include parent components beyond its own
1226 // name; its full name is then the dot-joined components. This is the single, converged
1227 // nesting rule — previously the Rust aggregate used `len > 1` while Python/Node also gated
1228 // on `name_field`, which could disagree.
1229 let nested = params.len() > 1;
1230 Some(AccessorSpec {
1231 singular: resource.descriptor.singular.clone(),
1232 params,
1233 nested,
1234 })
1235 }
1236
1237 pub(crate) fn methods(&self) -> impl Iterator<Item = MethodHandler<'_>> {
1238 self.plan.methods.iter().map(|plan| MethodHandler {
1239 plan,
1240 metadata: self.metadata,
1241 config: self.config,
1242 })
1243 }
1244
1245 /// The methods that have a `google.api.http` route — the ones the REST emitters (client,
1246 /// handler, server, aggregate) can generate. Routeless methods (no HTTP annotation) are emitted
1247 /// only by the ConnectRPC client, which iterates [`Self::methods`] directly.
1248 pub(crate) fn rest_methods(&self) -> impl Iterator<Item = MethodHandler<'_>> {
1249 self.methods().filter(|m| m.plan.has_http_route)
1250 }
1251
1252 /// The ergonomic, public-facing client type for this service.
1253 ///
1254 /// For a resource-scoped service this is the **scoped resource client** (e.g. `CatalogClient`,
1255 /// bound to a name) — the type callers and language bindings use. For a resource-less service it
1256 /// is the single flat client (e.g. `TagAssignmentClient`). The low-level transport client (one
1257 /// async method per RPC) is [`Self::low_level_client_type`].
1258 pub(crate) fn client_type(&self) -> Ident {
1259 if let Some(resource) = self.resource() {
1260 format_ident!(
1261 "{}",
1262 format!("{} client", resource.descriptor.singular).to_case(Case::Pascal)
1263 )
1264 } else {
1265 format_ident!(
1266 "{}Client",
1267 self.plan
1268 .service_name
1269 .trim_end_matches("Service")
1270 .trim_end_matches('s')
1271 )
1272 }
1273 }
1274
1275 /// The low-level transport client type (`{ client: CloudClient, base_url: Url }` with one async
1276 /// method per RPC), emitted by [`client`](super::client).
1277 ///
1278 /// For a resource-scoped service this is `<Singular>ServiceClient` (e.g. `CatalogServiceClient`),
1279 /// distinct from the ergonomic scoped [`Self::client_type`] (`CatalogClient`). For a
1280 /// resource-less service there is no scoped client, so the low-level client *is* the
1281 /// [`Self::client_type`] (e.g. `TagAssignmentClient`).
1282 pub(crate) fn low_level_client_type(&self, protocol: crate::codegen::ClientProtocol) -> Ident {
1283 // ConnectRPC: a single adapter type per service, named deterministically from the proto
1284 // service so it never collides with (nor is confused for) the connect-rust `{Service}Client`
1285 // it wraps. Connect has no scoped resource client, so this is always the low-level type.
1286 if protocol == crate::codegen::ClientProtocol::Connect {
1287 return format_ident!(
1288 "{}ConnectClient",
1289 self.plan.service_name.to_case(Case::Pascal)
1290 );
1291 }
1292 if let Some(resource) = self.resource() {
1293 format_ident!(
1294 "{}ServiceClient",
1295 resource.descriptor.singular.to_case(Case::Pascal)
1296 )
1297 } else {
1298 self.client_type()
1299 }
1300 }
1301
1302 /// The connect-rust generated client type for this service, as a fully-qualified path.
1303 ///
1304 /// Resolves to `<connect_client_path>::<ServiceName>Client` (e.g.
1305 /// `headwaters_proto::connect_gen::headwaters::read::v1::ReadServiceClient`). Used only in
1306 /// [`ClientProtocol::Connect`](crate::codegen::ClientProtocol) mode; errors if
1307 /// `connect_client_path` is unset.
1308 pub(crate) fn connect_inner_client_path(&self) -> crate::Result<syn::Path> {
1309 let base = self.config.connect_client_path.as_deref().ok_or_else(|| {
1310 crate::Error::Build(
1311 "client_protocol = Connect requires `connect_client_path` to be set".to_string(),
1312 )
1313 })?;
1314 let full = format!(
1315 "{base}::{}Client",
1316 self.plan.service_name.to_case(Case::Pascal)
1317 );
1318 syn::parse_str(&full).map_err(|source| crate::error::Error::InvalidRustPath {
1319 path: full.clone(),
1320 message: "connect_client_path + service client name".to_string(),
1321 source,
1322 })
1323 }
1324
1325 /// The ergonomic scoped resource client type, or `None` for a resource-less service.
1326 ///
1327 /// Same as [`Self::client_type`] but `Option`-typed to make "resource-scoped only" explicit at
1328 /// call sites that generate the scoped client / its accessor.
1329 pub(crate) fn scoped_client_type(&self) -> Option<Ident> {
1330 self.resource().map(|_| self.client_type())
1331 }
1332
1333 /// Whether this service's resource can be modeled by an ergonomic resource-scoped client.
1334 ///
1335 /// The scoped client (and the aggregate client's `<singular>(…)` accessor) captures a
1336 /// resource's identity as an ordered list of *string* name components and fills each method's
1337 /// path parameters from them. A resource whose scoped instance methods carry a non-string path
1338 /// parameter — e.g. a model version addressed by its parent's composite `{full_name}` plus an
1339 /// integer `{version}` — cannot be expressed that way, so no scoped client is emitted for it
1340 /// (its methods remain available on the low-level service client). Returns `false` in that case.
1341 pub(crate) fn supports_scoped_client(&self) -> bool {
1342 !self
1343 .methods()
1344 .filter(|m| m.is_scoped_instance_method())
1345 .any(|m| {
1346 m.required_parameters().any(|p| {
1347 p.is_path_param()
1348 && !matches!(
1349 p.field_type().base_type,
1350 crate::parsing::types::BaseType::String
1351 )
1352 })
1353 })
1354 }
1355
1356 /// The module segment under which this service's generated models live.
1357 ///
1358 /// Models are emitted per proto *package* (e.g. `unitycatalog.tags.v1` → module `tags`),
1359 /// which is not always the same as the service's `base_path` (derived from the service
1360 /// *name*, e.g. `TagPoliciesService` → `tag_policies`). Use the package's leaf-before-version
1361 /// segment so the models import path matches where the proto plugin actually wrote the types.
1362 /// Falls back to `base_path` when the package is empty or unparsable.
1363 fn models_segment(&self) -> String {
1364 package_models_segment(&self.plan.package)
1365 .map_or_else(|| self.plan.base_path.clone(), str::to_string)
1366 }
1367
1368 /// The version segment under which this service's generated models live; see
1369 /// [`package_models_version`].
1370 fn models_version(&self) -> String {
1371 package_models_version(&self.plan.package).to_string()
1372 }
1373
1374 pub(crate) fn models_path(&self) -> syn::Path {
1375 // Templates (and each service's substitution) are validated by `generate_code` before
1376 // any `ServiceHandler` is used, so skip the redundant re-validation `new` would do.
1377 ModelsPath::from_template(&self.config.models_path_template)
1378 .resolve(&self.models_segment(), &self.models_version())
1379 }
1380
1381 /// `use` globs for the models of other packages that this service's RPCs take or return.
1382 pub(crate) fn foreign_models_imports(&self) -> Vec<TokenStream> {
1383 let template = ModelsPath::from_template(&self.config.models_path_template);
1384 self.methods()
1385 .flat_map(|m| {
1386 [
1387 m.plan.metadata.input_type.as_str(),
1388 m.plan.metadata.output_type.as_str(),
1389 ]
1390 })
1391 .filter_map(message_package)
1392 .filter(|package| *package != self.plan.package && !package.starts_with("google."))
1393 .unique()
1394 .filter_map(|package| {
1395 let segment = package_models_segment(&package)?;
1396 let path = template.resolve(segment, package_models_version(&package));
1397 Some(quote! { use #path::*; })
1398 })
1399 .collect()
1400 }
1401
1402 pub(crate) fn models_path_crate(&self) -> syn::Path {
1403 ModelsPath::from_template(&self.config.models_path_crate_template)
1404 .resolve(&self.models_segment(), &self.models_version())
1405 }
1406
1407 /// The `#[cfg(feature = "<model package>")]` gate for this service's client items, or no
1408 /// tokens when [`RestClientOptions::resource_features`] is off.
1409 ///
1410 /// The feature is named after the model package segment (see [`Self::models_segment`]), so
1411 /// services sharing a package (e.g. two tag services over `tags.v1`) share a feature.
1412 pub(crate) fn client_feature_cfg(&self) -> TokenStream {
1413 if self.config.rest_client.resource_features {
1414 let feature = self.models_segment();
1415 quote! { #[cfg(feature = #feature)] }
1416 } else {
1417 quote! {}
1418 }
1419 }
1420
1421 /// The models *root* module path — the prefix of [`Self::models_path`] before the
1422 /// per-service/version segments (e.g. `example_common::models` for the template
1423 /// `example_common::models::{service}::v1`).
1424 ///
1425 /// The `Py*` wrapper `#[pyclass]` types are re-exported here (see
1426 /// `generate_models_mod`), so the Python bindings import them from this path.
1427 /// Returns `None` if the template has no `{service}` placeholder to split on.
1428 pub(crate) fn models_root_path(&self) -> Option<syn::Path> {
1429 let models_path = ModelsPath::from_template(&self.config.models_path_template);
1430 syn::parse_str(models_path.root()?).ok()
1431 }
1432}
1433
1434/// The package of a fully qualified message name (`.my.shares.v1.ShareInfo` → `my.shares.v1`):
1435/// the lowercase segments before the message names.
1436fn message_package(type_name: &str) -> Option<String> {
1437 let segments: Vec<&str> = type_name
1438 .trim_start_matches('.')
1439 .split('.')
1440 .take_while(|s| s.starts_with(|c: char| c.is_ascii_lowercase()))
1441 .collect();
1442 (!segments.is_empty()).then(|| segments.join("."))
1443}
1444
1445/// The models module segment of a proto package: the one before the version (`my.tags.v1` →
1446/// `tags`), else the last (`my.tags` → `tags`).
1447pub(crate) fn package_models_segment(package: &str) -> Option<&str> {
1448 let segs: Vec<&str> = package.split('.').filter(|s| !s.is_empty()).collect();
1449 match segs.as_slice() {
1450 [.., name, version] if is_version_segment(version) => Some(name),
1451 [.., name] => Some(name),
1452 [] => None,
1453 }
1454}
1455
1456/// The version segment of a proto package (`my.tags.v1` → `v1`, `my.agents.v0alpha1` →
1457/// `v0alpha1`); `v1` when it has none.
1458pub(crate) fn package_models_version(package: &str) -> &str {
1459 let segs: Vec<&str> = package.split('.').filter(|s| !s.is_empty()).collect();
1460 match segs.as_slice() {
1461 [.., _name, version] if is_version_segment(version) => version,
1462 _ => "v1",
1463 }
1464}
1465
1466/// Whether a proto package segment is a version marker like `v1`, `v2beta1`, etc.
1467fn is_version_segment(seg: &str) -> bool {
1468 seg.strip_prefix('v')
1469 .is_some_and(|rest| rest.starts_with(|c: char| c.is_ascii_digit()))
1470}
1471
1472/// The per-method counterpart of [`ServiceHandler`], scoped to a single RPC.
1473///
1474/// Pairs one analyzed [`MethodPlan`](crate::analysis::MethodPlan) with the same [`CodeGenMetadata`]
1475/// and [`CodeGenConfig`] borrows, and exposes the per-method derivations the generators rely on:
1476/// input/output message lookup, builder and binding method names, parameter plans
1477/// ([`Self::param_plan`]), and the language-agnostic [`EmitShape`](crate::analysis::EmitShape).
1478/// Yielded by [`ServiceHandler::methods`]; each generator iterates these to emit one unit of code
1479/// (a trait method, a client method, a builder, …) per RPC.
1480pub(crate) struct MethodHandler<'a> {
1481 plan: &'a MethodPlan,
1482 metadata: &'a CodeGenMetadata,
1483 config: &'a CodeGenConfig,
1484}
1485
1486impl MethodHandler<'_> {
1487 /// The query-string key of `param` (see [`QueryParam::key`]).
1488 pub(crate) fn query_key<'p>(
1489 &self,
1490 param: &'p crate::analysis::QueryParam,
1491 ) -> std::borrow::Cow<'p, str> {
1492 let input = self.plan.metadata.input_type.trim_start_matches('.');
1493 let json_names = self
1494 .config
1495 .rest_client
1496 .json_query_packages
1497 .iter()
1498 .any(|package| {
1499 input
1500 .strip_prefix(package.as_str())
1501 .is_some_and(|rest| rest.starts_with('.'))
1502 });
1503 param.key(json_names)
1504 }
1505
1506 pub(crate) fn is_collection_method(&self) -> bool {
1507 crate::analysis::is_collection_method(
1508 &self.plan.request_type,
1509 &self.plan.metadata.method_name,
1510 )
1511 }
1512
1513 /// The language-agnostic [`EmitShape`] for this method, derived once from its request type and
1514 /// name. The Python/Node/TypeScript emitters branch on this instead of each re-matching
1515 /// `RequestType`.
1516 pub(crate) fn emit_shape(&self) -> crate::analysis::EmitShape {
1517 crate::analysis::emit_shape(&self.plan.request_type, self.is_collection_method())
1518 }
1519
1520 /// Whether this method is emitted as an instance method on a resource-scoped client
1521 /// (`get`/`update`/`delete` or a resource-targeted custom `POST`/`PATCH`).
1522 ///
1523 /// This is the exact set the bindings previously matched in their `resource_client_method`
1524 /// dispatch: standard `Get`/`Update`/`Delete`, plus a non-collection custom `POST`/`PATCH`
1525 /// (e.g. `POST /catalogs/{name}:rotateToken`). Other custom verbs (a custom `GET`/`DELETE`) are
1526 /// deliberately **not** surfaced on the scoped client, matching prior behavior.
1527 pub(crate) fn is_scoped_instance_method(&self) -> bool {
1528 // Standard instance verbs are always scoped.
1529 if matches!(
1530 self.plan.request_type,
1531 RequestType::Get | RequestType::Update | RequestType::Delete
1532 ) {
1533 return true;
1534 }
1535 // A *resource-targeted* custom RPC (any verb) is a scoped instance method: it has at least
1536 // one path param (so it operates on a specific resource) and is not a collection method
1537 // (those — list / create / factory — live on the root client). This surfaces custom reads
1538 // like `GetTableExists` / `GetPermissions` (`GET …/{name}/exists`) and resource-targeted
1539 // custom POST/PATCH (`…/{name}:rotateToken`) on the scoped client, rather than leaving their
1540 // generated builders orphaned/unreachable.
1541 matches!(self.plan.request_type, RequestType::Custom(_))
1542 && !self.is_collection_method()
1543 && self.plan.path_parameters().next().is_some()
1544 }
1545
1546 /// The ordered binding parameter list for `mode`, with shared filtering applied once.
1547 ///
1548 /// Required params come before optional ones (matching the previous
1549 /// `required_parameters().chain(optional_parameters())` order). Filtering reproduces exactly what
1550 /// the old per-language `collection_method_parameters` / `resource_method_parameters` did,
1551 /// keyed on [`EmitShape`]:
1552 ///
1553 /// - **collection shapes** (`List`, `Create`): never drop path params (collection methods live
1554 /// on the root client and pass everything); drop `page_token` for `List` (streaming handles
1555 /// pagination).
1556 /// - **instance shapes** (`GetUpdate`, `Delete`): drop path params in [`BindingMode::Scoped`]
1557 /// (the scoped client already holds them); never filter `page_token`.
1558 ///
1559 /// Per-language emitters render the returned `&RequestParam`s (applying their own type mapping /
1560 /// capability filters) instead of each re-deriving this filtered list.
1561 pub(crate) fn param_plan(&self, mode: BindingMode) -> Vec<&RequestParam> {
1562 let (required, optional) = self.param_plan_split(mode);
1563 required.into_iter().chain(optional).collect()
1564 }
1565
1566 /// Like [`Self::param_plan`] but keeps the required/optional partition, for emitters (e.g.
1567 /// TypeScript) that render the two groups differently (required positional args vs an optional
1568 /// `options` object). Both groups have the same shared path/`page_token` filtering applied.
1569 pub(crate) fn param_plan_split(
1570 &self,
1571 mode: BindingMode,
1572 ) -> (Vec<&RequestParam>, Vec<&RequestParam>) {
1573 use crate::analysis::EmitShape;
1574 let shape = self.emit_shape();
1575 let is_collection_shape = matches!(shape, EmitShape::List | EmitShape::Create);
1576 let drop_path = !is_collection_shape && mode == BindingMode::Scoped;
1577 let drop_page_token = shape == EmitShape::List;
1578 let keep = |p: &&RequestParam| {
1579 let is_dropped_path = drop_path && p.is_path_param();
1580 let is_dropped_page_token = drop_page_token && p.name() == "page_token";
1581 !is_dropped_path && !is_dropped_page_token
1582 };
1583 let required = self.required_parameters().filter(keep).collect();
1584 let optional = self.optional_parameters().filter(keep).collect();
1585 (required, optional)
1586 }
1587
1588 pub(crate) fn output_message(&self) -> Option<MessageMeta<'_>> {
1589 if self.plan.metadata.output_type.ends_with("Empty") {
1590 return None;
1591 }
1592 self.metadata
1593 .get_message_meta(&self.plan.metadata.output_type)
1594 }
1595
1596 pub(crate) fn output_type(&self) -> Option<Ident> {
1597 self.output_message()
1598 .map(|t| extract_type_ident(&t.info.name))
1599 }
1600
1601 /// The type a client call resolves to: the output message, `bool` for an existence check,
1602 /// or `()`.
1603 pub(crate) fn result_type(&self) -> TokenStream {
1604 match self.output_type() {
1605 Some(output) => quote! { #output },
1606 None if self.is_existence_check() => quote! { bool },
1607 None => quote! { () },
1608 }
1609 }
1610
1611 /// A `HEAD` request without an output: it returns `bool`, `false` for a 404.
1612 pub(crate) fn is_existence_check(&self) -> bool {
1613 self.plan.http_method.eq_ignore_ascii_case("head") && self.output_type().is_none()
1614 }
1615
1616 pub(crate) fn list_output_field(&self) -> Option<&MessageField> {
1617 self.output_message()?
1618 .info
1619 .fields
1620 .iter()
1621 .find(|f| !f.name.contains("page_token"))
1622 }
1623
1624 pub(crate) fn input_message(&self) -> Option<MessageMeta<'_>> {
1625 if self.plan.metadata.input_type == "Empty" {
1626 return None;
1627 }
1628 self.metadata
1629 .get_message_meta(&self.plan.metadata.input_type)
1630 }
1631
1632 pub(crate) fn input_type(&self) -> Option<Ident> {
1633 self.input_message()
1634 .map(|t| extract_type_ident(&t.info.name))
1635 }
1636
1637 pub(crate) fn builder_type(&self) -> Ident {
1638 format_ident!("{}Builder", self.plan.metadata.method_name)
1639 }
1640
1641 /// The method name to expose in language bindings.
1642 ///
1643 /// Prefers the gnostic `operation_id` annotation (`gnostic.openapi.v3.operation`) when present,
1644 /// snake-cased, so proto authors can give bindings a hand-tuned name. Falls back to the
1645 /// proto-derived name ([`MethodPlan::base_method_ident`]).
1646 ///
1647 /// This affects only the *emitted* binding method name — the call into the underlying Rust
1648 /// client still uses [`MethodPlan::base_method_ident`], so the handler/client API is unchanged.
1649 pub(crate) fn binding_method_name(&self) -> Ident {
1650 format_ident!("{}", self.binding_method_name_str())
1651 }
1652
1653 /// The snake_case binding method name (see [`Self::binding_method_name`]).
1654 ///
1655 /// Provided as a `String` so the string-templated TypeScript generator can camel-case it and
1656 /// stay consistent with the snake→camel name NAPI derives for the native method.
1657 pub(crate) fn binding_method_name_str(&self) -> String {
1658 match self.plan.metadata.operation.as_ref() {
1659 Some(op) if !op.operation_id.is_empty() => op.operation_id.to_case(Case::Snake),
1660 _ => self.plan.handler_function_name.clone(),
1661 }
1662 }
1663
1664 /// Get type representation for rust depending on context
1665 ///
1666 /// Depending on context we may want concrete types (e.g. 'String') or more flexible types (e.g. 'Into<String d>')
1667 pub(crate) fn field_type(&self, field_type: &UnifiedType, ctx: RenderContext) -> syn::Type {
1668 let rust_type = types::unified_to_rust(field_type, ctx);
1669 // `rust_type` comes from this crate's own proto→Rust mapping, not user input, so an
1670 // unparsable result is a bug in `unified_to_rust`. Panic with the offending string so
1671 // the broken mapping is obvious (the generated-code parse test in
1672 // `tests/generation_integration.rs` is the broader guard for output validity).
1673 syn::parse_str(&rust_type).unwrap_or_else(|e| {
1674 panic!("internal: unified_to_rust produced an invalid type `{rust_type}`: {e}")
1675 })
1676 }
1677
1678 /// Get field assignment TokenStream for constructor
1679 pub(crate) fn field_assignment(
1680 &self,
1681 field_type: &UnifiedType,
1682 field_ident: &proc_macro2::Ident,
1683 ctx: &RenderContext,
1684 ) -> TokenStream {
1685 types::field_assignment(field_type, field_ident, ctx, self.config.runtime)
1686 }
1687
1688 pub(crate) fn required_parameters(&self) -> impl Iterator<Item = &RequestParam> {
1689 self.plan
1690 .parameters
1691 .iter()
1692 .filter(|param| !param.is_optional())
1693 }
1694
1695 pub(crate) fn optional_parameters(&self) -> impl Iterator<Item = &RequestParam> {
1696 self.plan
1697 .parameters
1698 .iter()
1699 .filter(|param| param.is_optional())
1700 }
1701
1702 /// Split body fields into required and optional subsets.
1703 pub(crate) fn split_body_fields(&self) -> (Vec<&BodyField>, Vec<&BodyField>) {
1704 split_body_fields(self.plan)
1705 }
1706}
1707
1708/// Extract the final type name from a fully qualified protobuf type and convert to Ident.
1709///
1710/// Thin `Ident` adapter over [`crate::utils::extract_simple_type_name`] — the single source of
1711/// truth for last-segment extraction. (Not [`crate::utils::extract_qualified_type_name`], which
1712/// parent-prefixes nested types — different behavior.)
1713pub(crate) fn extract_type_ident(full_type: &str) -> Ident {
1714 format_ident!("{}", crate::utils::extract_simple_type_name(full_type))
1715}
1716
1717#[cfg(test)]
1718mod tests {
1719 use super::message_package;
1720
1721 #[test]
1722 fn message_package_stops_at_the_message_name() {
1723 assert_eq!(
1724 message_package(".unitycatalog.shares.v1.ShareInfo").as_deref(),
1725 Some("unitycatalog.shares.v1")
1726 );
1727 assert_eq!(
1728 message_package("demo.v1.Outer.Inner").as_deref(),
1729 Some("demo.v1")
1730 );
1731 assert_eq!(message_package("Empty"), None);
1732 }
1733}