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 config;
45mod handler;
46pub(crate) mod node;
47mod python;
48mod resource_client;
49mod resources;
50mod server;
51mod tokens;
52mod wasm;
53
54pub use config::{
55    BindingsConfig, CodeGenConfig, CodeGenOutput, DEFAULT_TRANSPORT_TYPE_PATH, ModelsPath, Runtime,
56};
57pub(crate) use tokens::{doc_tokens, format_tokens, format_tokens_static};
58
59/// How a language binding lowers a method's call into the underlying Rust client.
60///
61/// Resource services expose a scoped client (e.g. `CatalogClient`) that captures path params;
62/// resource-less services have no such client, so their methods live on the root client and must
63/// pass every param themselves (the "flat" shape, treated like collection methods).
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65pub(crate) enum BindingMode {
66    /// Method lives on a resource-scoped client that already holds the path params.
67    Scoped,
68    /// Method lives on the root client and receives all params, including path params.
69    Flat,
70}
71
72/// Language-agnostic description of a service's resource accessor method on the aggregate client.
73///
74/// Computed once from the IR (`ServicePlan::resource_accessor_params`) and rendered by each language
75/// emitter, so the param list and nested-vs-flat decision are shared by construction.
76pub(crate) struct AccessorSpec {
77    /// The resource singular, used as the accessor method name (e.g. `catalog`).
78    pub(crate) singular: String,
79    /// The ordered accessor params (ancestors + `<singular>_name` leaf), e.g.
80    /// `["catalog_name", "schema_name"]` or `["name"]`.
81    pub(crate) params: Vec<String>,
82    /// Whether the resource is nested (has parent components); its full name is the dot-joined
83    /// params and the accessor delegates through a `<singular>_from_full_name` helper.
84    pub(crate) nested: bool,
85}
86
87impl AccessorSpec {
88    /// The `format!` template for joining the params into a composite full name
89    /// (e.g. `"{}.{}"` for two params). Only meaningful when [`Self::nested`] is true.
90    pub(crate) fn join_format(&self) -> String {
91        std::iter::repeat_n("{}", self.params.len())
92            .collect::<Vec<_>>()
93            .join(".")
94    }
95}
96
97impl MethodPlan {
98    pub(crate) fn resource_client_method(&self) -> Ident {
99        match &self.scoped_verb {
100            Some(verb) => format_ident!("{}", verb),
101            None => format_ident!("{}", self.handler_function_name),
102        }
103    }
104
105    pub(crate) fn base_method_ident(&self) -> Ident {
106        format_ident!("{}", self.handler_function_name)
107    }
108}
109
110/// Generate all code described by `config` from `metadata`.
111///
112/// Writes the following outputs, depending on which [`CodeGenOutput`] fields are `Some`:
113///
114/// | Field | Contents |
115/// |-------|----------|
116/// | `output.common` | Axum extractor code, per-service `mod.rs` (always written) |
117/// | `output.models_gen` | `labels.rs` resource-enum file (falls back to `common` if `None`) |
118/// | `output.server` | Handler trait + Axum route wiring per service |
119/// | `output.client` | HTTP client structs and request builders per service |
120/// | `output.python` | PyO3 binding wrappers + `.pyi` typings stub |
121/// | `output.node` | NAPI binding wrappers |
122/// | `output.node_ts` | TypeScript client (`client.ts`) |
123///
124/// # Required fields
125///
126/// - `output.common` is always required.
127/// - `bindings` must be `Some` when any of `output.python`, `output.node`, or `output.node_ts`
128///   is `Some`; otherwise returns [`Error::MissingBindingsConfig`].
129/// - `models_path_template` and `models_path_crate_template` must be valid Rust path templates
130///   (containing `{service}`); invalid templates return [`Error::InvalidModelsPathTemplate`].
131///
132/// # Optional fields
133///
134/// Setting `generate_resource_enum` to `false` skips `labels.rs` generation.
135/// Setting `bindings` to `None` skips all language-binding output.
136///
137/// Call [`CodeGenConfig::validate`] before this function to surface config errors at
138/// construction time rather than mid-generation.
139///
140/// # Errors
141///
142/// Returns:
143///
144/// - [`Error::InvalidModelsPathTemplate`] if `models_path_template` or `models_path_crate_template`
145///   is not a valid Rust path template.
146/// - [`Error::InvalidConfigPath`] if `context_type_path` or `result_type_path` is not a valid Rust
147///   path.
148/// - [`Error::MissingBindingsConfig`] if any of the `python`, `node`, or `node_ts` outputs is
149///   requested without a `bindings` config.
150/// - [`Error::MissingAnnotation`] or [`Error::Build`] surfaced while analyzing the metadata (e.g. a
151///   required `google.api` annotation is missing, or a service cannot be planned).
152/// - [`Error::InvalidRustPath`] if a proto message or field maps to an invalid Rust path, and
153///   [`Error::InvalidErrorTypePath`] if the configured `error_type_path` is not a valid Rust path.
154/// - [`Error::GeneratedParse`] if an internal generator produces tokens that do not parse as a Rust
155///   file.
156/// - [`Error::Io`] if creating the models directory or writing any generated file fails.
157pub fn generate_code(metadata: &CodeGenMetadata, config: &CodeGenConfig) -> Result<()> {
158    // Validate templates early so callers get a clean error before any generation starts.
159    let models_path = ModelsPath::new(&config.models_path_template)?;
160    let models_path_crate = ModelsPath::new(&config.models_path_crate_template)?;
161
162    // The context/result type paths come from user config and are parsed as `syn::Path` deep in
163    // the per-method generators (handler/client/server/builder). Validate them once here so a
164    // typo surfaces as a clean error instead of panicking mid-generation.
165    for (field, path) in [
166        ("context_type", &config.context_type_path),
167        ("result_type", &config.result_type_path),
168    ] {
169        syn::parse_str::<syn::Path>(path).map_err(|source| Error::InvalidConfigPath {
170            field,
171            path: path.clone(),
172            source,
173        })?;
174    }
175
176    // Validate that bindings config is present when language-binding output is requested.
177    if (config.output.python.is_some()
178        || config.output.node.is_some()
179        || config.output.node_ts.is_some())
180        && config.bindings.is_none()
181    {
182        return Err(Error::MissingBindingsConfig);
183    }
184
185    let plan = analyze_metadata(metadata)?;
186
187    // Validate every service's `{service}` substitution up front so a base path that isn't a
188    // valid Rust path segment fails cleanly here rather than panicking inside per-service codegen.
189    for service in &plan.services {
190        models_path.validate_for(&service.base_path)?;
191        models_path_crate.validate_for(&service.base_path)?;
192    }
193
194    let common_code = generate_common_code(&plan, metadata, config)?;
195    output::write_generated_code(&common_code, &config.output.common)?;
196
197    if config.output.models.is_some() {
198        let subdir = config
199            .output
200            .models_subdir_path()
201            .expect("models is Some so subdir is Some");
202        std::fs::create_dir_all(&subdir).map_err(Error::Io)?;
203
204        if config.generate_resource_enum {
205            let resource_enum = resources::generate_resource_enum(
206                &plan,
207                metadata,
208                config,
209                config.error_type_path.as_deref(),
210            )?;
211            let mut models_files = GeneratedCode {
212                files: std::collections::HashMap::new(),
213            };
214            models_files
215                .files
216                .insert("labels.rs".to_string(), resource_enum);
217            output::write_generated_code(&models_files, &subdir)?;
218        }
219
220        let gen_dir = config.models_gen_dir.as_deref().unwrap_or("../gen");
221        let include_labels = config.generate_resource_enum;
222        let mod_content = generate_models_mod(&plan.services, gen_dir, include_labels, metadata);
223        // Route through write_generated_code so the `// @generated` header is prepended like
224        // every other emitted file, rather than writing it raw.
225        let mut mod_files = GeneratedCode {
226            files: std::collections::HashMap::new(),
227        };
228        mod_files.files.insert("mod.rs".to_string(), mod_content);
229        output::write_generated_code(&mod_files, &subdir)?;
230    }
231
232    if let Some(ref server_dir) = config.output.server {
233        let server_code = generate_server_code(&plan, metadata, config)?;
234        output::write_generated_code(&server_code, server_dir)?;
235    }
236
237    if let Some(ref client_dir) = config.output.client {
238        let client_code = generate_client_code(&plan, metadata, config)?;
239        output::write_generated_code(&client_code, client_dir)?;
240    }
241
242    if let Some(ref python_dir) = config.output.python {
243        let python_code = generate_python_code(&plan, metadata, config)?;
244        output::write_generated_code(&python_code, python_dir)?;
245    }
246
247    if let Some(ref node_dir) = config.output.node {
248        let node_code = generate_node_code(&plan, metadata, config)?;
249        output::write_generated_code(&node_code, node_dir)?;
250    }
251
252    if let Some(ref node_ts_dir) = config.output.node_ts {
253        let node_ts_code = generate_node_ts_code(&plan, metadata, config)?;
254        output::write_generated_code(&node_ts_code, node_ts_dir)?;
255    }
256
257    if let Some(ref wasm_dir) = config.output.wasm {
258        let wasm_code = generate_wasm_code(&plan, metadata, config)?;
259        output::write_generated_code(&wasm_code, wasm_dir)?;
260    }
261
262    Ok(())
263}
264
265fn generate_common_code(
266    plan: &GenerationPlan,
267    metadata: &CodeGenMetadata,
268    config: &CodeGenConfig,
269) -> Result<GeneratedCode> {
270    let mut files = HashMap::new();
271
272    for service in &plan.services {
273        let handler = ServiceHandler {
274            plan: service,
275            metadata,
276            config,
277        };
278        let server_code = server::generate_common(&handler)?;
279        files.insert(format!("{}/server.rs", service.base_path), server_code);
280        let module_code = generate_common_module();
281        files.insert(format!("{}/mod.rs", service.base_path), module_code);
282    }
283
284    let module_code = main_module(&plan.services);
285    files.insert("mod.rs".to_string(), module_code);
286
287    Ok(GeneratedCode { files })
288}
289
290fn generate_server_code(
291    plan: &GenerationPlan,
292    metadata: &CodeGenMetadata,
293    config: &CodeGenConfig,
294) -> Result<GeneratedCode> {
295    let mut files = HashMap::new();
296
297    for service in &plan.services {
298        let handler = ServiceHandler {
299            plan: service,
300            metadata,
301            config,
302        };
303        let trait_code = handler::generate(&handler)?;
304        files.insert(format!("{}/handler.rs", service.base_path), trait_code);
305        let server_code = server::generate_server(&handler)?;
306        files.insert(format!("{}/server.rs", service.base_path), server_code);
307        let module_code = generate_server_module(service);
308        files.insert(format!("{}/mod.rs", service.base_path), module_code);
309    }
310
311    let module_code = main_module(&plan.services);
312    files.insert("mod.rs".to_string(), module_code);
313
314    Ok(GeneratedCode { files })
315}
316
317fn generate_python_code(
318    plan: &GenerationPlan,
319    metadata: &CodeGenMetadata,
320    config: &CodeGenConfig,
321) -> Result<GeneratedCode> {
322    let mut files = HashMap::new();
323
324    let handlers = plan
325        .services
326        .iter()
327        .map(|service| ServiceHandler {
328            plan: service,
329            metadata,
330            config,
331        })
332        .collect_vec();
333
334    // Per-service scoped client modules are only meaningful for resource-scoped services. The
335    // methods of resource-less services live on the root client (emitted by `main_module`), so
336    // they don't get a scoped module of their own.
337    for service in handlers.iter().filter(|s| s.is_resource_scoped()) {
338        let python_code = python::generate(service)?;
339        files.insert(format!("{}.rs", service.plan.base_path), python_code);
340    }
341
342    let module_code = python::main_module(&handlers)?;
343    files.insert("mod.rs".to_string(), module_code);
344
345    let python_typings_code = python::generate_typings(&handlers);
346    files.insert(
347        config.output.python_typings_filename.clone(),
348        python_typings_code,
349    );
350
351    Ok(GeneratedCode { files })
352}
353
354fn generate_node_code(
355    plan: &GenerationPlan,
356    metadata: &CodeGenMetadata,
357    config: &CodeGenConfig,
358) -> Result<GeneratedCode> {
359    let mut files = HashMap::new();
360
361    let handlers = plan
362        .services
363        .iter()
364        .map(|service| ServiceHandler {
365            plan: service,
366            metadata,
367            config,
368        })
369        .collect_vec();
370
371    // See `generate_python_code`: only resource-scoped services get a per-service scoped module;
372    // resource-less services' methods live on the root client.
373    for service in handlers.iter().filter(|s| s.is_resource_scoped()) {
374        let napi_code = node::generate(service)?;
375        files.insert(format!("{}.rs", service.plan.base_path), napi_code);
376    }
377
378    let module_code = node::main_module(&handlers)?;
379    files.insert("mod.rs".to_string(), module_code);
380
381    Ok(GeneratedCode { files })
382}
383
384fn generate_node_ts_code(
385    plan: &GenerationPlan,
386    metadata: &CodeGenMetadata,
387    config: &CodeGenConfig,
388) -> Result<GeneratedCode> {
389    let handlers = plan
390        .services
391        .iter()
392        .map(|service| ServiceHandler {
393            plan: service,
394            metadata,
395            config,
396        })
397        .collect_vec();
398
399    let ts_code = node::typescript::generate_client_ts(&handlers);
400    let mut files = HashMap::new();
401    files.insert("client.ts".to_string(), ts_code);
402
403    // The `models/index.ts` barrel re-exports every service's generated protobuf-es modules so
404    // `client.ts`'s `from "./models"` imports (message and service request/response types) resolve.
405    let models_barrel = node::typescript::generate_models_barrel(&handlers);
406    files.insert("models/index.ts".to_string(), models_barrel);
407
408    Ok(GeneratedCode { files })
409}
410
411/// Generate the WASM/browser `#[wasm_bindgen]` binding layer plus its `.d.ts`.
412///
413/// Emits `bindings.rs` (a `#[wasm_bindgen]` wrapper per service client, exchanging request/response
414/// values as plain JS objects via `serde-wasm-bindgen`) and `client.d.ts` (TypeScript declarations
415/// for JS/TS consumers). The wrapped clients are the generated low-level clients, which must be
416/// generated with the WASM transport (`transport_type_path = "olai_http_wasm::WasmClient"`).
417fn generate_wasm_code(
418    plan: &GenerationPlan,
419    metadata: &CodeGenMetadata,
420    config: &CodeGenConfig,
421) -> Result<GeneratedCode> {
422    let handlers = plan
423        .services
424        .iter()
425        .map(|service| ServiceHandler {
426            plan: service,
427            metadata,
428            config,
429        })
430        .collect_vec();
431
432    let mut files = HashMap::new();
433    files.insert(
434        "bindings.rs".to_string(),
435        wasm::generate_bindings(&handlers)?,
436    );
437    files.insert("client.d.ts".to_string(), wasm::generate_dts(&handlers));
438
439    Ok(GeneratedCode { files })
440}
441
442fn generate_client_code(
443    plan: &GenerationPlan,
444    metadata: &CodeGenMetadata,
445    config: &CodeGenConfig,
446) -> Result<GeneratedCode> {
447    let mut files = HashMap::new();
448
449    for service in &plan.services {
450        let handler = ServiceHandler {
451            plan: service,
452            metadata,
453            config,
454        };
455        let client_code = client::generate(&handler)?;
456        files.insert(format!("{}/client.rs", service.base_path), client_code);
457        let builder_code = builder::generate(&handler)?;
458        files.insert(format!("{}/builders.rs", service.base_path), builder_code);
459
460        // Ergonomic resource-scoped client, when enabled and the service manages a resource.
461        // `plan` is passed so the emitter can resolve child services for navigation/create accessors.
462        let has_resource_client = if config.output.generate_resource_clients {
463            resource_client::generate(&handler, plan)?
464                .map(|code| {
465                    files.insert(format!("{}/resource.rs", service.base_path), code);
466                })
467                .is_some()
468        } else {
469            false
470        };
471
472        let module_code = generate_client_module(has_resource_client);
473        files.insert(format!("{}/mod.rs", service.base_path), module_code);
474    }
475
476    // Top-level aggregate root client (e.g. `UnityCatalogClient`), emitted only when a bindings
477    // config (which carries the aggregate name) is present.
478    let has_aggregate = aggregate::generate(plan, metadata, config)?
479        .map(|aggregate_code| {
480            files.insert("client.rs".to_string(), aggregate_code);
481        })
482        .is_some();
483
484    let module_code = generate_client_main_module(&plan.services, has_aggregate);
485    files.insert("mod.rs".to_string(), module_code);
486
487    Ok(GeneratedCode { files })
488}
489
490fn generate_common_module() -> String {
491    let tokens = quote! {
492        #[cfg(feature = "axum")]
493        pub mod server;
494    };
495    format_tokens_static(tokens)
496}
497
498fn generate_server_module(service: &ServicePlan) -> String {
499    let handler_ident = format_ident!("{}", service.handler_name);
500    let tokens = quote! {
501        pub use handler::#handler_ident;
502
503        mod handler;
504        #[cfg(feature = "axum")]
505        pub mod server;
506    };
507    format_tokens_static(tokens)
508}
509
510fn generate_client_module(has_resource_client: bool) -> String {
511    // Emit the scoped resource client module only when it was generated for this service.
512    let resource_module = if has_resource_client {
513        quote! {
514            pub use resource::*;
515            pub mod resource;
516        }
517    } else {
518        quote! {}
519    };
520    let tokens = quote! {
521        pub use client::*;
522        pub use builders::*;
523        #resource_module
524
525        pub mod client;
526        pub mod builders;
527    };
528    format_tokens_static(tokens)
529}
530
531pub fn main_module(services: &[ServicePlan]) -> String {
532    let service_modules: Vec<TokenStream> = services
533        .iter()
534        .map(|s| {
535            let module_name = format_ident!("{}", s.base_path);
536            quote! { pub mod #module_name; }
537        })
538        .collect();
539
540    let tokens = quote! {
541        #(#service_modules)*
542    };
543    format_tokens_static(tokens)
544}
545
546fn generate_client_main_module(services: &[ServicePlan], has_aggregate: bool) -> String {
547    let service_modules: Vec<TokenStream> = services
548        .iter()
549        .map(|s| {
550            let module_name = format_ident!("{}", s.base_path);
551            quote! { pub mod #module_name; }
552        })
553        .collect();
554
555    // Export the aggregate root client (`client.rs`) when it was generated. Per-service files live
556    // under `<base_path>/`, so a top-level `client` module never collides with them.
557    let aggregate_module = if has_aggregate {
558        quote! {
559            pub mod client;
560            pub use client::*;
561        }
562    } else {
563        quote! {}
564    };
565
566    let tokens = quote! {
567        #(#service_modules)*
568
569        #aggregate_module
570
571        use futures::Future;
572
573        pub(super) fn stream_paginated<F, Fut, S, T>(
574            state: S,
575            op: F,
576        ) -> impl futures::Stream<Item = crate::Result<T>>
577        where
578            F: Fn(S, Option<String>) -> Fut + Copy,
579            Fut: Future<Output = crate::Result<(T, S, Option<String>)>>,
580        {
581            enum PaginationState<T> {
582                Start(T),
583                HasMore(T, String),
584                Done,
585            }
586
587            futures::stream::unfold(
588                PaginationState::Start(state),
589                move |state| async move {
590                    let (s, page_token) = match state {
591                        PaginationState::Start(s) => (s, None),
592                        PaginationState::HasMore(s, page_token) if !page_token.is_empty() => {
593                            (s, Some(page_token))
594                        }
595                        _ => {
596                            return None;
597                        }
598                    };
599
600                    let (resp, s, continuation) = match op(s, page_token).await {
601                        Ok(resp) => resp,
602                        Err(e) => return Some((Err(e), PaginationState::Done)),
603                    };
604
605                    let next_state = match continuation {
606                        Some(token) => PaginationState::HasMore(s, token),
607                        None => PaginationState::Done,
608                    };
609
610                    Some((Ok(resp), next_state))
611                },
612            )
613        }
614    };
615    format_tokens_static(tokens)
616}
617
618/// Generate the `mod.rs` for `crates/common/src/models/`.
619///
620/// Emits `pub mod <service> { pub mod <version> { include!(...) } }` blocks for every
621/// service in the plan, plus static re-exports and module declarations.
622///
623/// `gen_dir` is the relative path (from the subdir) to the prost-generated files,
624/// e.g. `"../gen"`.
625///
626/// When `include_labels` is `true`, also emits `pub mod labels; pub use labels::{ObjectLabel, Resource};`.
627///
628/// `metadata` is used to discover all resource-annotated messages (even those not returned
629/// directly by an RPC) so they can be included in `pub use` re-exports.
630pub fn generate_models_mod(
631    services: &[ServicePlan],
632    gen_dir: &str,
633    include_labels: bool,
634    metadata: &CodeGenMetadata,
635) -> String {
636    // Services are keyed by proto package, not by service: multiple services can share a
637    // package (e.g. `CatalogService` and `SchemaService` both in `example.catalog.v1`), and
638    // emitting one `pub mod`/`pub use` block per service would produce duplicate modules and
639    // imports (E0428). Group by package first, then emit once per package.
640    //
641    // package = "unitycatalog.catalogs.v1"
642    // parts   = ["unitycatalog", "catalogs", "v1"]
643    // service module = second-to-last segment ("catalogs"); version = last segment ("v1").
644    let mod_segments = |svc: &ServicePlan| -> (String, String) {
645        let parts: Vec<&str> = svc.package.split('.').collect();
646        if parts.len() >= 2 {
647            (
648                parts[parts.len() - 2].to_string(),
649                parts[parts.len() - 1].to_string(),
650            )
651        } else {
652            (svc.base_path.clone(), "v1".to_string())
653        }
654    };
655
656    // One entry per package, in sorted order, with the union of re-exported type names.
657    let mut packages: BTreeMap<String, (String, String, BTreeSet<String>)> = BTreeMap::new();
658    for svc in services {
659        let (svc_seg, ver_seg) = mod_segments(svc);
660        let entry = packages
661            .entry(svc.package.clone())
662            .or_insert_with(|| (svc_seg, ver_seg, BTreeSet::new()));
663        let type_names = &mut entry.2;
664
665        // Re-export managed resources for this service...
666        type_names.extend(svc.managed_resources.iter().map(|r| r.type_name.clone()));
667
668        // ...plus every resource-annotated message in the same package (catches nested types
669        // like `Column` that aren't direct RPC return types but still carry annotations).
670        let fq_prefix = format!(".{}.", svc.package);
671        let fq_pkg = format!(".{}", svc.package);
672        for (fq_name, msg_info) in &metadata.messages {
673            if msg_info.resource_descriptor.is_some()
674                && (fq_name.starts_with(&fq_prefix) || fq_name.starts_with(&fq_pkg))
675            {
676                let simple = fq_name
677                    .rfind('.')
678                    .map(|i| &fq_name[i + 1..])
679                    .unwrap_or(fq_name.as_str());
680                type_names.insert(simple.to_string());
681            }
682        }
683    }
684
685    let mut service_mods: Vec<TokenStream> = Vec::new();
686    let mut reexports: Vec<TokenStream> = Vec::new();
687    for (package, (svc_seg, ver_seg, type_names)) in &packages {
688        let svc_mod = format_ident!("{}", svc_seg);
689        let ver_mod = format_ident!("{}", ver_seg);
690
691        let main_include = format!("./{}/{}.rs", gen_dir, package);
692        let tonic_include = format!("./{}/{}.tonic.rs", gen_dir, package);
693        service_mods.push(quote! {
694            pub mod #svc_mod {
695                pub mod #ver_mod {
696                    include!(#main_include);
697                    #[cfg(feature = "grpc")]
698                    include!(#tonic_include);
699                }
700            }
701        });
702
703        for type_name in type_names {
704            let type_ident = format_ident!("{}", type_name);
705            reexports.push(quote! {
706                pub use #svc_mod::#ver_mod::#type_ident;
707            });
708        }
709    }
710
711    let labels_decl: TokenStream = if include_labels {
712        quote! {
713            pub mod labels;
714            pub use labels::{ObjectLabel, Resource};
715        }
716    } else {
717        quote! {}
718    };
719
720    let tokens = quote! {
721        use std::collections::HashMap;
722
723        #labels_decl
724
725        #(#reexports)*
726
727        pub type PropertyMap = HashMap<String, serde_json::Value>;
728
729        #(#service_mods)*
730    };
731
732    format_tokens_static(tokens)
733}
734
735/// An in-memory bundle of generated source files awaiting output.
736///
737/// Each generator phase (common, server, client, bindings, …) returns one of these; the
738/// [`crate::output`] writer then materializes the [`files`](Self::files) onto disk, prepending the
739/// `// @generated` header to each. Keeping generation and writing separate lets generation stay
740/// pure (no I/O) and makes the produced files easy to assert on in tests.
741#[derive(Debug)]
742pub struct GeneratedCode {
743    /// Generated files keyed by path relative to the phase's output directory.
744    pub files: HashMap<String, String>,
745}
746
747impl CodeGenMetadata {
748    fn get_message_meta(&self, message_name: &str) -> Option<MessageMeta<'_>> {
749        self.messages.get(message_name).map(|info| MessageMeta {
750            info,
751            metadata: self,
752        })
753    }
754}
755
756/// A protobuf message paired with the metadata needed to resolve its references.
757///
758/// Wraps a single [`MessageInfo`](crate::parsing::MessageInfo) together with a borrow of the whole
759/// [`CodeGenMetadata`] it came from, so the generators can look up a message's fields while still
760/// being able to resolve types that point at other messages in the same descriptor set. Obtained
761/// via `CodeGenMetadata::get_message_meta`.
762pub(crate) struct MessageMeta<'a> {
763    info: &'a MessageInfo,
764    #[allow(dead_code)]
765    metadata: &'a CodeGenMetadata,
766}
767
768/// The per-service generation context shared by every code generator.
769///
770/// Bundles the three borrows each generator phase needs for one service: its analyzed
771/// [`ServicePlan`](crate::analysis::ServicePlan), the full [`CodeGenMetadata`] (for resolving
772/// types referenced from other services), and the active [`CodeGenConfig`]. The various
773/// `codegen::*` submodules take a `&ServiceHandler` and read these plus the derived accessors on
774/// this type (client type names, models paths, [`Self::methods`]) rather than threading the three
775/// borrows separately. Methods are exposed as [`MethodHandler`]s via [`Self::methods`].
776pub(crate) struct ServiceHandler<'a> {
777    pub(crate) plan: &'a ServicePlan,
778    pub(crate) metadata: &'a CodeGenMetadata,
779    pub(crate) config: &'a CodeGenConfig,
780}
781
782impl ServiceHandler<'_> {
783    pub(crate) fn resource(&self) -> Option<&ManagedResource> {
784        self.plan.managed_resources.first()
785    }
786
787    /// Whether this service manages a `google.api.resource` and therefore gets a resource-scoped
788    /// client (e.g. `CatalogClient` bound to a name).
789    ///
790    /// Resource-less services (entity tag assignments, temporary credentials, delta commits) have
791    /// no scoped client; their methods are emitted on the root client and pass every param
792    /// directly ([`BindingMode::Flat`]). See [`Self::binding_mode`].
793    pub(crate) fn is_resource_scoped(&self) -> bool {
794        self.resource().is_some()
795    }
796
797    /// The [`BindingMode`] used to lower this service's methods in language bindings.
798    pub(crate) fn binding_mode(&self) -> BindingMode {
799        if self.is_resource_scoped() {
800            BindingMode::Scoped
801        } else {
802            BindingMode::Flat
803        }
804    }
805
806    /// The language-agnostic spec for this service's resource accessor on the aggregate client, or
807    /// `None` for resource-less services.
808    ///
809    /// All four accessor emitters (Rust aggregate, Python, Node, TypeScript) render from this single
810    /// spec rather than re-deriving the param list and nesting decision, which previously diverged
811    /// between the Rust aggregate and the bindings.
812    pub(crate) fn accessor_spec(&self) -> Option<AccessorSpec> {
813        let resource = self.resource()?;
814        let params = self.plan.resource_accessor_params.clone()?;
815        // A resource is nested when its accessor params include parent components beyond its own
816        // name; its full name is then the dot-joined components. This is the single, converged
817        // nesting rule — previously the Rust aggregate used `len > 1` while Python/Node also gated
818        // on `name_field`, which could disagree.
819        let nested = params.len() > 1;
820        Some(AccessorSpec {
821            singular: resource.descriptor.singular.clone(),
822            params,
823            nested,
824        })
825    }
826
827    pub(crate) fn methods(&self) -> impl Iterator<Item = MethodHandler<'_>> {
828        self.plan.methods.iter().map(|plan| MethodHandler {
829            plan,
830            metadata: self.metadata,
831            config: self.config,
832        })
833    }
834
835    /// The ergonomic, public-facing client type for this service.
836    ///
837    /// For a resource-scoped service this is the **scoped resource client** (e.g. `CatalogClient`,
838    /// bound to a name) — the type callers and language bindings use. For a resource-less service it
839    /// is the single flat client (e.g. `TagAssignmentClient`). The low-level transport client (one
840    /// async method per RPC) is [`Self::low_level_client_type`].
841    pub(crate) fn client_type(&self) -> Ident {
842        if let Some(resource) = self.resource() {
843            format_ident!(
844                "{}",
845                format!("{} client", resource.descriptor.singular).to_case(Case::Pascal)
846            )
847        } else {
848            format_ident!(
849                "{}Client",
850                self.plan
851                    .service_name
852                    .trim_end_matches("Service")
853                    .trim_end_matches('s')
854            )
855        }
856    }
857
858    /// The low-level transport client type (`{ client: CloudClient, base_url: Url }` with one async
859    /// method per RPC), emitted by [`client`](super::client).
860    ///
861    /// For a resource-scoped service this is `<Singular>ServiceClient` (e.g. `CatalogServiceClient`),
862    /// distinct from the ergonomic scoped [`Self::client_type`] (`CatalogClient`). For a
863    /// resource-less service there is no scoped client, so the low-level client *is* the
864    /// [`Self::client_type`] (e.g. `TagAssignmentClient`).
865    pub(crate) fn low_level_client_type(&self) -> Ident {
866        if let Some(resource) = self.resource() {
867            format_ident!(
868                "{}ServiceClient",
869                resource.descriptor.singular.to_case(Case::Pascal)
870            )
871        } else {
872            self.client_type()
873        }
874    }
875
876    /// The ergonomic scoped resource client type, or `None` for a resource-less service.
877    ///
878    /// Same as [`Self::client_type`] but `Option`-typed to make "resource-scoped only" explicit at
879    /// call sites that generate the scoped client / its accessor.
880    pub(crate) fn scoped_client_type(&self) -> Option<Ident> {
881        self.resource().map(|_| self.client_type())
882    }
883
884    /// The module segment under which this service's generated models live.
885    ///
886    /// Models are emitted per proto *package* (e.g. `unitycatalog.tags.v1` → module `tags`),
887    /// which is not always the same as the service's `base_path` (derived from the service
888    /// *name*, e.g. `TagPoliciesService` → `tag_policies`). Use the package's leaf-before-version
889    /// segment so the models import path matches where the proto plugin actually wrote the types.
890    /// Falls back to `base_path` when the package is empty or unparsable.
891    fn models_segment(&self) -> String {
892        let segs: Vec<&str> = self
893            .plan
894            .package
895            .split('.')
896            .filter(|s| !s.is_empty())
897            .collect();
898        match segs.as_slice() {
899            // `unitycatalog.tags.v1` → `tags`
900            [.., name, version] if is_version_segment(version) => name.to_string(),
901            // `unitycatalog.tags` → `tags`
902            [.., name] => name.to_string(),
903            // No package info: fall back to the service-name-derived base path.
904            _ => self.plan.base_path.clone(),
905        }
906    }
907
908    pub(crate) fn models_path(&self) -> syn::Path {
909        // Templates (and each service's substitution) are validated by `generate_code` before
910        // any `ServiceHandler` is used, so skip the redundant re-validation `new` would do.
911        ModelsPath::from_template(&self.config.models_path_template).resolve(&self.models_segment())
912    }
913
914    pub(crate) fn models_path_crate(&self) -> syn::Path {
915        ModelsPath::from_template(&self.config.models_path_crate_template)
916            .resolve(&self.models_segment())
917    }
918}
919
920/// Whether a proto package segment is a version marker like `v1`, `v2beta1`, etc.
921fn is_version_segment(seg: &str) -> bool {
922    seg.strip_prefix('v')
923        .is_some_and(|rest| rest.starts_with(|c: char| c.is_ascii_digit()))
924}
925
926/// The per-method counterpart of [`ServiceHandler`], scoped to a single RPC.
927///
928/// Pairs one analyzed [`MethodPlan`](crate::analysis::MethodPlan) with the same [`CodeGenMetadata`]
929/// and [`CodeGenConfig`] borrows, and exposes the per-method derivations the generators rely on:
930/// input/output message lookup, builder and binding method names, parameter plans
931/// ([`Self::param_plan`]), and the language-agnostic [`EmitShape`](crate::analysis::EmitShape).
932/// Yielded by [`ServiceHandler::methods`]; each generator iterates these to emit one unit of code
933/// (a trait method, a client method, a builder, …) per RPC.
934pub(crate) struct MethodHandler<'a> {
935    plan: &'a MethodPlan,
936    metadata: &'a CodeGenMetadata,
937    config: &'a CodeGenConfig,
938}
939
940impl MethodHandler<'_> {
941    pub(crate) fn is_collection_method(&self) -> bool {
942        crate::analysis::is_collection_method(
943            &self.plan.request_type,
944            &self.plan.metadata.method_name,
945        )
946    }
947
948    /// The language-agnostic [`EmitShape`] for this method, derived once from its request type and
949    /// name. The Python/Node/TypeScript emitters branch on this instead of each re-matching
950    /// `RequestType`.
951    pub(crate) fn emit_shape(&self) -> crate::analysis::EmitShape {
952        crate::analysis::emit_shape(&self.plan.request_type, self.is_collection_method())
953    }
954
955    /// Whether this method is emitted as an instance method on a resource-scoped client
956    /// (`get`/`update`/`delete` or a resource-targeted custom `POST`/`PATCH`).
957    ///
958    /// This is the exact set the bindings previously matched in their `resource_client_method`
959    /// dispatch: standard `Get`/`Update`/`Delete`, plus a non-collection custom `POST`/`PATCH`
960    /// (e.g. `POST /catalogs/{name}:rotateToken`). Other custom verbs (a custom `GET`/`DELETE`) are
961    /// deliberately **not** surfaced on the scoped client, matching prior behavior.
962    pub(crate) fn is_scoped_instance_method(&self) -> bool {
963        // Standard instance verbs are always scoped.
964        if matches!(
965            self.plan.request_type,
966            RequestType::Get | RequestType::Update | RequestType::Delete
967        ) {
968            return true;
969        }
970        // A *resource-targeted* custom RPC (any verb) is a scoped instance method: it has at least
971        // one path param (so it operates on a specific resource) and is not a collection method
972        // (those — list / create / factory — live on the root client). This surfaces custom reads
973        // like `GetTableExists` / `GetPermissions` (`GET …/{name}/exists`) and resource-targeted
974        // custom POST/PATCH (`…/{name}:rotateToken`) on the scoped client, rather than leaving their
975        // generated builders orphaned/unreachable.
976        matches!(self.plan.request_type, RequestType::Custom(_))
977            && !self.is_collection_method()
978            && self.plan.path_parameters().next().is_some()
979    }
980
981    /// The ordered binding parameter list for `mode`, with shared filtering applied once.
982    ///
983    /// Required params come before optional ones (matching the previous
984    /// `required_parameters().chain(optional_parameters())` order). Filtering reproduces exactly what
985    /// the old per-language `collection_method_parameters` / `resource_method_parameters` did,
986    /// keyed on [`EmitShape`]:
987    ///
988    /// - **collection shapes** (`List`, `Create`): never drop path params (collection methods live
989    ///   on the root client and pass everything); drop `page_token` for `List` (streaming handles
990    ///   pagination).
991    /// - **instance shapes** (`GetUpdate`, `Delete`): drop path params in [`BindingMode::Scoped`]
992    ///   (the scoped client already holds them); never filter `page_token`.
993    ///
994    /// Per-language emitters render the returned `&RequestParam`s (applying their own type mapping /
995    /// capability filters) instead of each re-deriving this filtered list.
996    pub(crate) fn param_plan(&self, mode: BindingMode) -> Vec<&RequestParam> {
997        let (required, optional) = self.param_plan_split(mode);
998        required.into_iter().chain(optional).collect()
999    }
1000
1001    /// Like [`Self::param_plan`] but keeps the required/optional partition, for emitters (e.g.
1002    /// TypeScript) that render the two groups differently (required positional args vs an optional
1003    /// `options` object). Both groups have the same shared path/`page_token` filtering applied.
1004    pub(crate) fn param_plan_split(
1005        &self,
1006        mode: BindingMode,
1007    ) -> (Vec<&RequestParam>, Vec<&RequestParam>) {
1008        use crate::analysis::EmitShape;
1009        let shape = self.emit_shape();
1010        let is_collection_shape = matches!(shape, EmitShape::List | EmitShape::Create);
1011        let drop_path = !is_collection_shape && mode == BindingMode::Scoped;
1012        let drop_page_token = shape == EmitShape::List;
1013        let keep = |p: &&RequestParam| {
1014            let is_dropped_path = drop_path && p.is_path_param();
1015            let is_dropped_page_token = drop_page_token && p.name() == "page_token";
1016            !is_dropped_path && !is_dropped_page_token
1017        };
1018        let required = self.required_parameters().filter(keep).collect();
1019        let optional = self.optional_parameters().filter(keep).collect();
1020        (required, optional)
1021    }
1022
1023    pub(crate) fn output_message(&self) -> Option<MessageMeta<'_>> {
1024        if self.plan.metadata.output_type.ends_with("Empty") {
1025            return None;
1026        }
1027        self.metadata
1028            .get_message_meta(&self.plan.metadata.output_type)
1029    }
1030
1031    pub(crate) fn output_type(&self) -> Option<Ident> {
1032        self.output_message()
1033            .map(|t| extract_type_ident(&t.info.name))
1034    }
1035
1036    /// The method's output type as a `syn::Type`, falling back to the unit type `()` when the
1037    /// RPC returns `Empty` (i.e. [`Self::output_message`] is `None`).
1038    ///
1039    /// Binding emitters splice the response type into `Result<T>` wrappers
1040    /// (e.g. `PyUnityCatalogResult<T>`); interpolating a bare `Option::None` would produce an
1041    /// uncompilable `PyUnityCatalogResult<>`. Use this so `Empty`-returning RPCs yield
1042    /// `PyUnityCatalogResult<()>` instead.
1043    pub(crate) fn output_type_or_unit(&self) -> syn::Type {
1044        match self.output_type() {
1045            Some(ident) => syn::parse_quote!(#ident),
1046            None => syn::parse_quote!(()),
1047        }
1048    }
1049
1050    pub(crate) fn list_output_field(&self) -> Option<&MessageField> {
1051        self.output_message()?
1052            .info
1053            .fields
1054            .iter()
1055            .find(|f| !f.name.contains("page_token"))
1056    }
1057
1058    pub(crate) fn input_message(&self) -> Option<MessageMeta<'_>> {
1059        if self.plan.metadata.input_type == "Empty" {
1060            return None;
1061        }
1062        self.metadata
1063            .get_message_meta(&self.plan.metadata.input_type)
1064    }
1065
1066    pub(crate) fn input_type(&self) -> Option<Ident> {
1067        self.input_message()
1068            .map(|t| extract_type_ident(&t.info.name))
1069    }
1070
1071    pub(crate) fn builder_type(&self) -> Ident {
1072        format_ident!("{}Builder", self.plan.metadata.method_name)
1073    }
1074
1075    /// The method name to expose in language bindings.
1076    ///
1077    /// Prefers the gnostic `operation_id` annotation (`gnostic.openapi.v3.operation`) when present,
1078    /// snake-cased, so proto authors can give bindings a hand-tuned name. Falls back to the
1079    /// proto-derived name ([`MethodPlan::base_method_ident`]).
1080    ///
1081    /// This affects only the *emitted* binding method name — the call into the underlying Rust
1082    /// client still uses [`MethodPlan::base_method_ident`], so the handler/client API is unchanged.
1083    pub(crate) fn binding_method_name(&self) -> Ident {
1084        format_ident!("{}", self.binding_method_name_str())
1085    }
1086
1087    /// The snake_case binding method name (see [`Self::binding_method_name`]).
1088    ///
1089    /// Provided as a `String` so the string-templated TypeScript generator can camel-case it and
1090    /// stay consistent with the snake→camel name NAPI derives for the native method.
1091    pub(crate) fn binding_method_name_str(&self) -> String {
1092        match self.plan.metadata.operation.as_ref() {
1093            Some(op) if !op.operation_id.is_empty() => op.operation_id.to_case(Case::Snake),
1094            _ => self.plan.handler_function_name.clone(),
1095        }
1096    }
1097
1098    /// Get type representation for rust depending on context
1099    ///
1100    /// Depending on context we may want concrete types (e.g. 'String') or more flexible types (e.g. 'Into<String d>')
1101    pub(crate) fn field_type(&self, field_type: &UnifiedType, ctx: RenderContext) -> syn::Type {
1102        let rust_type = types::unified_to_rust(field_type, ctx);
1103        // `rust_type` comes from this crate's own proto→Rust mapping, not user input, so an
1104        // unparsable result is a bug in `unified_to_rust`. Panic with the offending string so
1105        // the broken mapping is obvious (the generated-code parse test in
1106        // `tests/generation_integration.rs` is the broader guard for output validity).
1107        syn::parse_str(&rust_type).unwrap_or_else(|e| {
1108            panic!("internal: unified_to_rust produced an invalid type `{rust_type}`: {e}")
1109        })
1110    }
1111
1112    /// Get field assignment TokenStream for constructor
1113    pub(crate) fn field_assignment(
1114        &self,
1115        field_type: &UnifiedType,
1116        field_ident: &proc_macro2::Ident,
1117        ctx: &RenderContext,
1118    ) -> TokenStream {
1119        types::field_assignment(field_type, field_ident, ctx, self.config.runtime)
1120    }
1121
1122    pub(crate) fn required_parameters(&self) -> impl Iterator<Item = &RequestParam> {
1123        self.plan
1124            .parameters
1125            .iter()
1126            .filter(|param| !param.is_optional())
1127    }
1128
1129    pub(crate) fn optional_parameters(&self) -> impl Iterator<Item = &RequestParam> {
1130        self.plan
1131            .parameters
1132            .iter()
1133            .filter(|param| param.is_optional())
1134    }
1135
1136    /// Split body fields into required and optional subsets.
1137    pub(crate) fn split_body_fields(&self) -> (Vec<&BodyField>, Vec<&BodyField>) {
1138        split_body_fields(self.plan)
1139    }
1140}
1141
1142/// Extract the final type name from a fully qualified protobuf type and convert to Ident.
1143///
1144/// Thin `Ident` adapter over [`crate::utils::extract_simple_type_name`] — the single source of
1145/// truth for last-segment extraction. (Not [`crate::utils::extract_qualified_type_name`], which
1146/// parent-prefixes nested types — different behavior.)
1147pub(crate) fn extract_type_ident(full_type: &str) -> Ident {
1148    format_ident!("{}", crate::utils::extract_simple_type_name(full_type))
1149}