Skip to main content

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}