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