olai-codegen

Proto-driven code generator that transforms compiled protobuf descriptors into idiomatic Rust server and client code, resource registries, and optional Python/Node.js/TypeScript bindings.

Overview

The single source of truth is your .proto files. Annotate them with standard Google API extensions and run trestle generate to emit:

Pipeline

.proto files
    │
    ▼  buf build --as-file-descriptor-set
descriptor binary (.bin)
    │
    ▼  trestle generate
Source files (Rust, Python, TypeScript, …)

Proto Annotations

The generator reads standard Google API extensions. No custom annotations are required.

Example proto fragment (see proto/ for full working examples):

message Catalog {
  option (google.api.resource) = {
    type: "example.io/Catalog"
    pattern: "catalogs/{catalog}"
    singular: "catalog"
    plural: "catalogs"
  };
  string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
  string comment = 2 [(google.api.field_behavior) = OPTIONAL];
}

service CatalogService {
  rpc GetCatalog(GetCatalogRequest) returns (Catalog) {
    option (google.api.http) = { get: "/catalogs/{name}" };
  }
  rpc CreateCatalog(CreateCatalogRequest) returns (Catalog) {
    option (google.api.http) = { post: "/catalogs" body: "*" };
  }
}

Usage

Installation

olai-codegen is a library; code generation is driven by the trestle CLI. Install it from the workspace:

cargo install --path crates/trestle
trestle generate --help

Compile descriptors

buf build --as-file-descriptor-set -o api.bin

Config file (recommended)

A YAML config file is the preferred way to invoke trestle generate, particularly in multi-crate workspaces where output directories span crate boundaries. CLI flags override any value from the file. New projects scaffolded with trestle new use trestle.yaml (the older proto-gen.yaml name is still accepted).

# trestle.yaml
descriptors: api.bin
buf_gen: buf.gen.yaml  # optional: auto-derives models_path_template from prost plugin output

generate:
  output_common:  crates/server/src/gen/common
  output_models:  crates/common/src/models   # labels.rs lands in <output_models>/_gen/
  output_server:  crates/server/src/gen/server
  output_client:  crates/client/src/gen/client

  context_type: my_crate::RequestContext
  result_type:  my_crate::Result

  # Required when buf_gen is not set:
  models_path_template:       my_models::models::{service}::v1
  models_path_crate_template: crate::models::{service}::v1

  # Resource enum generation (requires olai-store)
  generate_resource_enum:      true
  generate_store_integration:  true
  generate_object_conversions: true
  error_type_path: crate::Error

  # Optional: language bindings
  python:
    output: crates/python/src/gen
    error_type: my_crate::PyError
    result_type: my_crate::PyResult

  typescript:
    output: crates/node/src/gen
    aggregate_client_name: MyClient
    client_crate_name: my_client

trestle generate --config trestle.yaml

Any field can be overridden at the command line:

trestle generate --config trestle.yaml --descriptors path/to/other.bin

CLI flags (all options)

trestle generate

trestle enrich-openapi

Merges proto-derived validation rules into an OpenAPI YAML spec.

The same top-level config file covers both subcommands:

descriptors: api.bin

enrich_openapi:
  spec: openapi/openapi.yaml
  jsonschema_dir: openapi/jsonschema
  camel_case: true

Recommended Project Structure

The structure described below is what trestle new databricks-app-rust produces. Use the CLI to scaffold a fresh project in this shape and skip the manual setup:

trestle new my-app --template databricks-app-rust --profile dbx-emulator

The rest of this section documents the same layout for users who'd rather wire things up by hand.

trestle generate writes generated files into output directories that can span multiple crates. The recommended layout separates concerns across four crates:

my-workspace/
├── Cargo.toml
├── proto/
│   └── my_service.proto          # source of truth
├── api.bin                       # buf build output
├── trestle.yaml
└── crates/
    ├── common/                   # shared model types
    │   └── src/
    │       └── models/
    │           ├── mod.rs        # hand-written: re-exports from _gen
    │           └── _gen/
    │               ├── labels.rs # generated: Resource enum, ObjectLabel, store impls
    │               └── mod.rs    # generated: include! prost output + re-exports
    │
    ├── server/                   # Axum service implementation
    │   └── src/
    │       ├── api/
    │       │   ├── context.rs    # hand-written: RequestContext (impl FromRequestParts)
    │       │   └── error.rs      # hand-written: Error type, Result alias
    │       └── gen/
    │           ├── common/       # generated: Axum extractors
    │           └── server/       # generated: handler traits + route handlers
    │
    └── client/                   # HTTP client
        └── src/
            └── gen/
                ├── common/       # generated: shared extractor types
                └── client/       # generated: client struct + request builders

What you write vs. what is generated

Hand-written (once, then stable):

• RequestContext — extracted from each Axum request (auth, request ID, etc.); must implement axum::extract::FromRequestParts
• Error / Result<T> — your service error type and result alias; used in all generated handler and client signatures
• parse_error_response — called by generated client code on non-2xx responses; signature: pub async fn parse_error_response(response: reqwest::Response) -> Error
• common/src/models/mod.rs — re-exports from _gen so the rest of the codebase imports from a stable path

Generated (overwritten on each run, do not edit):

• common/src/models/_gen/labels.rs — Resource enum, ObjectLabel enum, olai_store::Label impl, RESOURCE_DESCRIPTORS
• common/src/models/_gen/mod.rs — include! directives into prost output + pub use re-exports
• server/src/gen/common/ — Axum extractor structs for path/query/body parameters
• server/src/gen/server/ — *Handler trait (one async fn per RPC) + Axum route handler functions
• client/src/gen/client/ — *Client struct + with_* builder methods

Implementing a handler

The generated handler trait has one method per RPC. Implement it in your server crate:

use async_trait::async_trait;
use gen::server::catalogs::CatalogHandler;

struct MyCatalogHandler { /* db pool, etc. */ }

#[async_trait]
impl CatalogHandler for MyCatalogHandler {
    async fn get_catalog(
        &self,
        request: GetCatalogRequest,
        context: RequestContext,
    ) -> Result<Catalog> {
        // ...
    }
}

The generator emits a free Axum handler function per RPC (e.g. get_catalog, create_catalog) but does not build the Router for you — you wire the routes by hand. This is deliberate: only your app knows the URL prefixes to nest under, and a service often mixes generated CRUD handlers with hand-written ones (custom streaming responses, bespoke headers) that the generator can't see. Write one small router fn per service:

use axum::routing::{delete, get, patch, post};
use gen::server::catalogs::server::*;

fn catalogs_router<T, Cx>(handler: T) -> axum::Router
where
    T: CatalogHandler<Cx> + Clone + Send + Sync + 'static,
    Cx: axum::extract::FromRequestParts<T> + Send + 'static,
{
    axum::Router::new()
        .route("/catalogs", get(list_catalogs::<T, Cx>).post(create_catalog::<T, Cx>))
        .route(
            "/catalogs/{name}",
            get(get_catalog::<T, Cx>)
                .patch(update_catalog::<T, Cx>)
                .delete(delete_catalog::<T, Cx>),
        )
        .with_state(handler)
}

Then merge the per-service routers and nest them under your API prefix:

let handler = MyCatalogHandler::new();
let app = axum::Router::new()
    .nest("/api/2.1/unity-catalog", catalogs_router(handler.clone()).merge(schemas_router(handler)));

Using the generated client

use olai_http::CloudClient;
use gen::client::catalogs::CatalogClient;

let client = CatalogClient::new(
    CloudClient::new(/* credentials */),
    url::Url::parse("https://api.example.com/")?,
);

let catalog = client.get_catalog(&GetCatalogRequest { name: "my-catalog".into() }).await?;

Cargo.toml dependencies

Add these to the crates that consume the generated code:

# server/Cargo.toml
[dependencies]
axum = "0.8"
axum-extra = "0.10"
async-trait = "0.1"
olai-store = { path = "../../crates/olai-store" }  # if store integration enabled

# client/Cargo.toml
[dependencies]
olai-http = { path = "../../crates/olai-http" }
url = "2"
serde_json = "1"

# common/Cargo.toml — additional deps when generate_object_conversions = true
[dependencies]
olai-store = { path = "../../crates/olai-store" }
uuid = { version = "1", features = ["v4"] }
chrono = "0.4"
strum = { version = "0.26", features = ["derive"] }

Library Usage

trestle generate handles the full workflow for most use cases. If you need to embed generation into a Rust program (e.g. a custom build tool), the library API exposes the same three-step pipeline:

use olai_codegen::{
    parse_file_descriptor_set, generate_code,
    CodeGenConfig, CodeGenOutput,
};
use protobuf::descriptor::FileDescriptorSet;
use protobuf::Message;

let descriptor_bytes = std::fs::read("api.bin")?;
let fds = FileDescriptorSet::parse_from_bytes(&descriptor_bytes)?;
let metadata = parse_file_descriptor_set(&fds)?;

let config = CodeGenConfig {
    context_type_path: "crate::RequestContext".into(),
    result_type_path: "crate::Result".into(),
    models_path_template: "my_models::models::{service}::v1".into(),
    models_path_crate_template: "crate::models::{service}::v1".into(),
    output: CodeGenOutput {
        common: "src/gen/common".into(),
        server: Some("src/gen/server".into()),
        client: Some("src/gen/client".into()),
        ..Default::default()
    },
    ..Default::default()
};

generate_code(&metadata, &config)?;

Configuration Reference

CodeGenConfig

WASM / browser clients

Generated clients are decoupled from their HTTP transport: the client struct stores whatever type transport_type_path names, and the method bodies only use a small surface — per-verb builders (get/post/put/patch/delete) returning a builder with .json(..)/.query(..)/ .send(), whose response has .status()/.bytes(). [olai_http::CloudClient] (the default) and [olai_http_wasm::WasmClient] both satisfy this.

Set the transport to the WASM client to emit a wasm32-unknown-unknown-buildable client for browser front-ends. In trestle.yaml:

transport: wasm        # friendly alias; or set transport_type_path: "olai_http_wasm::WasmClient"

The WASM client carries no request signing, ring, tokio, or cloud-credential discovery — in the browser the session is managed by the browser (fetch is asked to include credentials, so cookies / auth headers ride along on same-origin or CORS-with-credentials requests). The cloud-only aggregate constructors (new_unauthenticated, new_with_token) are not emitted for a WASM transport; construct the client with new(transport, base_url). Pair with runtime: buffa for serde-native, dependency-light models.

JS/TS bindings (#[wasm_bindgen] + .d.ts)

Set output.wasm (or wasm.output in trestle.yaml) to also emit a #[wasm_bindgen] wrapper layer for JS/browser consumption. Two files are generated:

• bindings.rs — a #[wasm_bindgen] struct per service plus an aggregate root constructed from a base URL. Each RPC is an async method that takes/returns plain JS objects, marshalled with serde-wasm-bindgen (no Buffer/protobuf-es decode step). Compile this crate to wasm32-unknown-unknown and package with wasm-pack.
• client.d.ts — TypeScript declarations: the aggregate class (new ExampleClient(baseUrl)), its per-service accessors, and each method typed method(request: ReqType): Promise<RespType> with request/response types imported from ./models.

runtime: buffa          # required: serde-native models cross the JS boundary as objects
transport: wasm
wasm:
  output: ./web/src/gen  # emits bindings.rs + client.d.ts

// after `wasm-pack build`
import { ExampleClient } from "./gen";
const client = new ExampleClient(window.location.origin);
const cat = await client.catalog().createCatalog({ name: "c1" });  // browser attaches the session

The generated client crate must depend on wasm-bindgen, wasm-bindgen-futures, serde-wasm-bindgen, and olai-http-wasm.

CodeGenOutput

Examples

The proto/ directory contains fully-annotated example protobuf definitions and a pre-compiled descriptor (example.bin). Integration tests in tests/ run the full pipeline against these examples and serve as executable documentation of the generated output.