openapi_lambda_codegen/
lib.rs

1#![doc = include_str!("../README.md")]
2#![allow(clippy::too_many_arguments)]
3#![warn(missing_docs)]
4
5use crate::api::operation::collect_operations;
6
7use indexmap::IndexMap;
8use itertools::Itertools;
9use openapiv3::{OpenAPI, Operation};
10use proc_macro2::{Ident, Span, TokenStream};
11use quote::quote;
12use serde_json::json;
13use syn::parse2;
14
15use std::borrow::Cow;
16use std::collections::HashMap;
17use std::fs::File;
18use std::io::Write;
19use std::path::{Path, PathBuf};
20use std::process::Command;
21
22mod api;
23mod apigw;
24mod inline;
25mod model;
26mod reference;
27
28// Re-export since `Operation` is part of the public API (for filters), and that includes references
29// to other `openapiv3` types.
30pub use openapiv3;
31
32/// Cache of parsed OpenAPI documents.
33type DocCache = HashMap<PathBuf, serde_yaml::Mapping>;
34
35#[derive(Debug)]
36enum LambdaArnImpl {
37  /// Use a `!Sub` AWS CloudFormation intrinsic to resolve the Lambda ARN at deploy time.
38  ///
39  /// This should be used only if the OpenAPI spec will be embedded in the `DefinitionBody` of an
40  /// `AWS::Serverless::Api` resource or the `Body` of an `AWS::ApiGateway::RestApi` resource. Note
41  /// that in both cases, an `Fn::Transform` intrinsic with the `AWS::Include` transform is needed
42  /// to resolve `!Sub` intrinsics in the OpenAPI template. Otherwise, the template will be deployed
43  /// verbatim without substituting the Lambda ARN, which the API Gateway service will reject.
44  CloudFormation {
45    /// Logical ID of the Lambda function within the CloudFormation/SAM template (e.g.,
46    /// `PetstoreFunction.Arn`, `PetstoreFunction.Alias`, or `PetstoreFunctionAliasLive`).
47    ///
48    /// This logical ID is used for resolving the Lambda function's ARN at deploy time using
49    /// CloudFormation intrinsics (e.g., `Sub`). This way, the generated OpenAPI spec passed to
50    /// the `AWS::Serverless::Api` resource's `DefinitionBody` can be generic and support multiple
51    /// deployments in distinct environments.
52    logical_id: String,
53  },
54  /// Use a known ARN that can be provided directly to API Gateway or an infrastructure-as-code
55  /// (IaC) solution other than AWS CloudFormation/SAM.
56  Known {
57    api_gateway_region: String,
58    account_id: String,
59    function_region: String,
60    function_name: String,
61    alias_or_version: Option<String>,
62  },
63}
64
65impl LambdaArnImpl {
66  pub fn apigw_invocation_arn(&self) -> serde_json::Value {
67    match self {
68      LambdaArnImpl::CloudFormation { logical_id } => {
69        json!({
70          "Fn::Sub": format!(
71            "arn:aws:apigateway:${{AWS::Region}}:lambda:path/2015-03-31/functions/${{{logical_id}}}\
72             /invocations",
73          )
74        })
75      }
76      LambdaArnImpl::Known {
77        api_gateway_region,
78        account_id,
79        function_region,
80        function_name,
81        alias_or_version,
82      } => serde_json::Value::String(format!(
83        "arn:aws:apigateway:{api_gateway_region}:lambda:path/2015-03-31/functions/arn:aws\
84         :lambda:{function_region}:{account_id}:function:{function_name}{}/invocations",
85        alias_or_version
86          .as_ref()
87          .map(|alias| Cow::Owned(format!(":{alias}")))
88          .unwrap_or(Cow::Borrowed(""))
89      )),
90    }
91  }
92}
93
94/// Amazon Resource Name (ARN) for an AWS Lambda function.
95///
96/// This type represents an ARN either using variables (e.g., an AWS CloudFormation logical ID
97/// constructed via the [`cloud_formation`](LambdaArn::cloud_formation) method) or as a
98/// fully-resolved ARN via the [`known`](LambdaArn::known) method. It is used to populate the
99/// [`x-amazon-apigateway-integration`](https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-swagger-extensions-integration.html)
100/// OpenAPI extensions that Amazon API Gateway uses to determine which Lambda function should handle
101/// each API endpoint.
102#[derive(Debug)]
103pub struct LambdaArn(LambdaArnImpl);
104
105impl LambdaArn {
106  /// Construct a variable ARN that references an AWS CloudFormation or Serverless Application Model
107  /// (SAM) logical ID.
108  ///
109  /// The logical ID should reference one of the following resource types defined in your
110  /// CloudFormation/SAM template:
111  ///  * [`AWS::Serverless::Function`](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html)
112  ///  * [`AWS::Lambda::Function`](https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-lambda-function.html)
113  ///  * [`AWS::Lambda::Alias`](https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-lambda-alias.html)
114  ///    (e.g., by appending `.Alias` to the logical ID when specifying an
115  ///    [`AutoPublishAlias`](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html#sam-function-autopublishalias)
116  ///    on the `AWS::Serverless::Function` resource)
117  ///  * [`AWS::Lambda::Version`](https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-lambda-version.html)
118  ///    (e.g., by appending `.Version` to the logical ID when specifying an
119  ///    [`AutoPublishAlias`](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html#sam-function-autopublishalias)
120  ///    on the `AWS::Serverless::Function` resource)
121  ///
122  /// When using this method, be sure to include the `openapi-apigw.yaml` file in your
123  /// CloudFormation/SAM template with the `AWS::Include` transform. Otherwise, the variables will
124  /// not be substituted during deployment, and deployment will fail. For example (where
125  /// `.openapi-lambda` is the `out_dir` passed to [`CodeGenerator::new`]):
126  /// ```yaml
127  /// Resources:
128  ///   MyApi:
129  ///     Type: AWS::Serverless::Api
130  ///     Properties:
131  ///       Name: my-api
132  ///       StageName: prod
133  ///       DefinitionBody:
134  ///         Fn::Transform:
135  ///           Name: AWS::Include
136  ///           Parameters:
137  ///             Location: .openapi-lambda/openapi-apigw.yaml
138  /// ```
139  ///
140  /// # Example
141  ///
142  /// ```rust
143  /// # use openapi_lambda_codegen::LambdaArn;
144  /// # let _ =
145  /// LambdaArn::cloud_formation("MyApiFunction.Alias")
146  /// # ;
147  /// ```
148  pub fn cloud_formation<L>(logical_id: L) -> Self
149  where
150    L: Into<String>,
151  {
152    Self(LambdaArnImpl::CloudFormation {
153      logical_id: logical_id.into(),
154    })
155  }
156
157  /// Construct a fully-resolved AWS Lambda function ARN.
158  ///
159  /// The resulting ARN does not depend on any CloudFormation variables and is compatible with any
160  /// deployment method.
161  ///
162  /// # Arguments
163  ///
164  /// * `api_gateway_region` - Region containing the Amazon API Gateway (e.g., `us-east-1`)
165  /// * `account_id` - AWS account containing the AWS Lambda function
166  /// * `function_region` - Region containing the AWS Lambda function (e.g., `us-east-1`)
167  /// * `function_name` - Name of the AWS Lambda function
168  /// * `alias_or_version` - Optional Lambda function
169  ///   [version](https://docs.aws.amazon.com/lambda/latest/dg/configuration-versions.html) or
170  ///   [alias](https://docs.aws.amazon.com/lambda/latest/dg/configuration-aliases.html)
171  ///
172  /// # Example
173  ///
174  /// ```rust
175  /// # use openapi_lambda_codegen::LambdaArn;
176  /// # let _ =
177  /// LambdaArn::known(
178  ///   "us-east-1",
179  ///   "1234567890",
180  ///   "us-east-1",
181  ///   "my-api-function",
182  ///   Some("live".to_string()),
183  /// )
184  /// # ;
185  /// ```
186  pub fn known<A, F, G, R>(
187    api_gateway_region: G,
188    account_id: A,
189    function_region: R,
190    function_name: F,
191    alias_or_version: Option<String>,
192  ) -> Self
193  where
194    A: Into<String>,
195    F: Into<String>,
196    G: Into<String>,
197    R: Into<String>,
198  {
199    Self(LambdaArnImpl::Known {
200      api_gateway_region: api_gateway_region.into(),
201      account_id: account_id.into(),
202      function_region: function_region.into(),
203      function_name: function_name.into(),
204      alias_or_version,
205    })
206  }
207}
208
209type OpFilter = Box<dyn Fn(&Operation) -> bool + 'static>;
210
211/// Builder for generating code for a single API Lambda function.
212///
213/// An `ApiLambda` instance represents a collection of API endpoints handled by a single
214/// Lambda function. This could include all endpoints defined in an OpenAPI spec (i.e., a
215/// "mono-Lambda"), a single API endpoint, or a subset of the API. Larger Lambda binaries incur a
216/// greater
217/// [cold start](https://docs.aws.amazon.com/lambda/latest/operatorguide/execution-environments.html#cold-start-latency)
218/// cost than smaller binaries, so the granularity of API Lambda functions presents a tradeoff
219/// between performance and implementation/deployment complexity (i.e., more Lambda functions to
220/// manage).
221///
222/// Use the [`with_op_filter`](ApiLambda::with_op_filter) method to specify a closure that
223/// associates API endpoints with the corresponding Lambda function.
224///
225/// # Example
226///
227/// ```rust
228/// # use openapi_lambda_codegen::{ApiLambda, LambdaArn};
229/// # let _ =
230/// ApiLambda::new("backend", LambdaArn::cloud_formation("BackendApiFunction.Alias"))
231/// # ;
232/// ```
233pub struct ApiLambda {
234  mod_name: String,
235  lambda_arn: LambdaArnImpl,
236  op_filter: Option<OpFilter>,
237}
238
239impl ApiLambda {
240  /// Construct a new `ApiLambda`.
241  ///
242  /// # Arguments
243  ///
244  /// * `mod_name` - Name of the Rust module to generate (must be a valid Rust identifier)
245  /// * `lambda_arn` - Amazon Resource Name (ARN) of the AWS Lambda function that will handle
246  ///   requests to the corresponding API endpoints via Amazon API Gateway (see [`LambdaArn`])
247  pub fn new<M>(mod_name: M, lambda_arn: LambdaArn) -> Self
248  where
249    M: Into<String>,
250  {
251    Self {
252      lambda_arn: lambda_arn.0,
253      mod_name: mod_name.into(),
254      op_filter: None,
255    }
256  }
257
258  /// Define a filter to associate a subset of API endpoints with this Lambda function.
259  ///
260  /// Use this method when *not* implementing a "mono-Lambda" that handles all API endpoints. By
261  /// default, all API endpoints will be included unless this method is called.
262  ///
263  /// # Arguments
264  ///
265  /// * `op_filter` - Closure that returns `true` or `false` to indicate whether the given OpenAPI
266  ///   [`Operation`] (endpoint) will be handled by the corresponding Lambda function
267  ///
268  /// # Example
269  ///
270  /// ```rust
271  /// # use openapi_lambda_codegen::{ApiLambda, LambdaArn};
272  /// # let _ =
273  /// ApiLambda::new("backend", LambdaArn::cloud_formation("BackendApiFunction.Alias"))
274  ///   // Only include API endpoints with the `pet` tag.
275  ///   .with_op_filter(|op| op.tags.iter().any(|tag| tag == "pet"))
276  /// # ;
277  /// ```
278  pub fn with_op_filter<F>(mut self, op_filter: F) -> Self
279  where
280    F: Fn(&Operation) -> bool + 'static,
281  {
282    self.op_filter = Some(Box::new(op_filter));
283    self
284  }
285}
286
287/// OpenAPI Lambda code generator.
288///
289/// This code generator is intended to be called from a `build.rs` Rust
290/// [build script](https://doc.rust-lang.org/cargo/reference/build-scripts.html). It emits an
291/// `out.rs` file to the directory referenced by the `OUT_DIR` environment variable set by Cargo.
292/// This file defines a module named `models` containing Rust types for the input parameters and
293/// request/response bodies defined in the OpenAPI definition. It also defines one
294/// module for each call to [`add_api_lambda`](CodeGenerator::add_api_lambda), which defines an
295/// `Api` trait with one method for each operation (path + HTTP method) defined in the OpenAPI
296/// definition.
297///
298/// In addition, the generator writes the following files to the `out_dir` directory specified in
299/// the call to [`new`](CodeGenerator::new):
300///  * `openapi-apigw.yaml` - OpenAPI definition annotated with
301///    [`x-amazon-apigateway-integration`](https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-swagger-extensions-integration.html)
302///    extensions to be used by Amazon API Gateway. This file is also modified from the input
303///    OpenAPI definition to help adhere to the
304///    [subset of OpenAPI features](https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-known-issues.html#api-gateway-known-issues-rest-apis)
305///    supported by Amazon API Gateway. In particular, all references are merged into a single file,
306///    and `discriminator` properties are removed.
307///  * One file for each call to [`add_api_lambda`](CodeGenerator::add_api_lambda) named
308///    `<MODULE_NAME>_handler.rs`, where `<MODULE_NAME>` is the `mod_name` in the [`ApiLambda`]
309///    passed to `add_api_lambda`. This file contains a placeholder implementation of the
310///    corresponding `Api` trait. To get started, copy this file into `src/`, define a corresponding
311///    module (`<MODULE_NAME>_handler`) in `src/lib.rs`, and replace each instance of `todo!()` in
312///    the trait implementation.
313///
314/// # Examples
315///
316/// ## Mono-Lambda
317///
318/// The following invocation in `build.rs` uses a single Lambda function to handle all API endpoints:
319/// ```rust,no_run
320/// # use openapi_lambda_codegen::{ApiLambda, CodeGenerator, LambdaArn};
321/// CodeGenerator::new("openapi.yaml", ".openapi-lambda")
322///   .add_api_lambda(
323///     ApiLambda::new("backend", LambdaArn::cloud_formation("BackendApiFunction.Alias"))
324///   )
325///   .generate();
326/// ```
327///
328/// ## Multiple Lambda functions
329///
330/// The following invocation in `build.rs` uses multiple Lambda functions, each handling a subset of
331/// API endpoints:
332/// ```rust,no_run
333/// # use openapi_lambda_codegen::{ApiLambda, CodeGenerator, LambdaArn};
334/// CodeGenerator::new("openapi.yaml", ".openapi-lambda")
335///   .add_api_lambda(
336///     ApiLambda::new("pet", LambdaArn::cloud_formation("PetApiFunction.Alias"))
337///     // Only include API endpoints with the `pet` tag.
338///     .with_op_filter(|op| op.tags.iter().any(|tag| tag == "pet"))
339///   )
340///   .add_api_lambda(
341///     ApiLambda::new("store", LambdaArn::cloud_formation("StoreApiFunction.Alias"))
342///     // Only include API endpoints with the `store` tag.
343///     .with_op_filter(|op| op.tags.iter().any(|tag| tag == "store"))
344///   )
345///   .generate();
346/// ```
347pub struct CodeGenerator {
348  api_lambdas: IndexMap<String, ApiLambda>,
349  openapi_path: PathBuf,
350  out_dir: PathBuf,
351}
352
353impl CodeGenerator {
354  /// Construct a new `CodeGenerator`.
355  ///
356  /// # Arguments
357  ///
358  /// * `openapi_path` - Input path to OpenAPI definition in YAML format
359  /// * `out_dir` - Output directory path in which `openapi-apigw.yaml` and one
360  ///   `<MODULE_NAME>_handler.rs` file for each call to
361  ///    [`add_api_lambda`](CodeGenerator::add_api_lambda) will be written
362  pub fn new<P, O>(openapi_path: P, out_dir: O) -> Self
363  where
364    P: Into<PathBuf>,
365    O: Into<PathBuf>,
366  {
367    Self {
368      api_lambdas: IndexMap::new(),
369      openapi_path: openapi_path.into(),
370      out_dir: out_dir.into(),
371    }
372  }
373
374  /// Register an API Lambda function for code generation.
375  ///
376  /// Each call to this method will result in a module being generated that contains an `Api` trait
377  /// with methods for the corresponding API endpoints. See [`ApiLambda`] for further details.
378  pub fn add_api_lambda(mut self, builder: ApiLambda) -> Self {
379    if self.api_lambdas.contains_key(&builder.mod_name) {
380      panic!(
381        "API Lambda module names must be unique: found duplicate `{}`",
382        builder.mod_name
383      )
384    }
385
386    self.api_lambdas.insert(builder.mod_name.clone(), builder);
387    self
388  }
389
390  /// Emit generated code.
391  pub fn generate(self) {
392    let cargo_out_dir = std::env::var("OUT_DIR").expect("OUT_DIR env not set");
393    log::info!("writing Rust codegen to {cargo_out_dir}");
394    log::info!("writing OpenAPI codegen to {}", self.out_dir.display());
395
396    if !self.out_dir.exists() {
397      std::fs::create_dir_all(&self.out_dir).unwrap_or_else(|err| {
398        panic!(
399          "failed to create directory `{}`: {err}",
400          self.out_dir.display()
401        )
402      });
403    }
404
405    let openapi_file = File::open(&self.openapi_path)
406      .unwrap_or_else(|err| panic!("failed to open {}: {err}", self.openapi_path.display()));
407
408    let openapi_yaml: serde_yaml::Mapping =
409      serde_path_to_error::deserialize(serde_yaml::Deserializer::from_reader(&openapi_file))
410        .unwrap_or_else(|err| panic!("Failed to parse OpenAPI spec as YAML: {err}"));
411
412    let mut cached_external_docs = DocCache::new();
413
414    // Clippy in 1.70.0 raises a false positive here.
415    #[allow(clippy::redundant_clone)]
416    cached_external_docs.insert(self.openapi_path.to_path_buf(), openapi_yaml.clone());
417
418    println!("cargo:rerun-if-changed={}", self.openapi_path.display());
419
420    let openapi: OpenAPI =
421      serde_path_to_error::deserialize(serde_yaml::Value::Mapping(openapi_yaml))
422        .unwrap_or_else(|err| panic!("Failed to parse OpenAPI spec: {err}"));
423
424    let crate_import = self.crate_use_name();
425
426    // Merge any references to other OpenAPI files into the root OpenAPI definition, and replace
427    // any unnamed schemas that require named models to represent in Rust (e.g., enums) with named
428    // schemas in components.schemas. This simplifies the rest of the code generation process since
429    // we don't have to visit other files or worry about conflicting schema names.
430    let (openapi_inline, models) =
431      self.generate_models(self.inline_openapi(openapi, cached_external_docs));
432
433    let openapi_inline_mapping =
434      serde_path_to_error::serialize(&*openapi_inline, serde_yaml::value::Serializer)
435        .expect("failed to serialize OpenAPI spec");
436    let serde_yaml::Value::Mapping(openapi_inline_mapping) = openapi_inline_mapping else {
437      panic!("OpenAPI spec should be a mapping: {:#?}", &*openapi_inline);
438    };
439
440    let operations = collect_operations(&openapi_inline, &openapi_inline_mapping);
441    let operations_by_api_lambda = self
442      .api_lambdas
443      .values()
444      .flat_map(|api_lambda| {
445        operations
446          .iter()
447          .filter(|op| {
448            api_lambda
449              .op_filter
450              .as_ref()
451              .map(|op_filter| (*op_filter)(&op.op))
452              .unwrap_or(true)
453          })
454          .map(|op| (&api_lambda.mod_name, op))
455      })
456      .into_group_map();
457
458    operations_by_api_lambda
459      .iter()
460      .flat_map(|(mod_name, ops)| {
461        ops
462          .iter()
463          .map(|op| ((&op.method, &op.request_path), *mod_name))
464      })
465      .into_group_map()
466      .into_iter()
467      .for_each(|((method, request_path), mod_names)| {
468        if mod_names.len() > 1 {
469          panic!(
470            "endpoint {method} {request_path} is mapped to multiple API Lambdas: {mod_names:?}"
471          );
472        }
473      });
474
475    let operation_id_to_api_lambda = operations_by_api_lambda
476      .iter()
477      .flat_map(|(mod_name, ops)| {
478        ops.iter().map(|op| {
479          (
480            op.op
481              .operation_id
482              .as_ref()
483              .unwrap_or_else(|| panic!("no operation_id for {} {}", op.method, op.request_path))
484              .as_str(),
485            self
486              .api_lambdas
487              .get(*mod_name)
488              .expect("mod name should exist in api_lambdas"),
489          )
490        })
491      })
492      .collect::<HashMap<_, _>>();
493
494    let components_schemas = openapi_inline
495      .components
496      .as_ref()
497      .map(|components| Cow::Borrowed(&components.schemas))
498      .unwrap_or_else(|| Cow::Owned(IndexMap::new()));
499    let apis_out = operations_by_api_lambda
500      .iter()
501      .sorted_by_key(|(mod_name, _)| **mod_name)
502      .map(|(mod_name, ops)| {
503        self.gen_api_module(
504          mod_name,
505          ops,
506          &openapi_inline_mapping,
507          &components_schemas,
508          &models,
509        )
510      })
511      .collect::<TokenStream>();
512
513    self.gen_openapi_apigw(openapi_inline, &operation_id_to_api_lambda);
514
515    let models_out = models
516      .into_iter()
517      .sorted_by(|(ident_a, _), (ident_b, _)| ident_a.cmp(ident_b))
518      .map(|(_, model)| model)
519      .collect::<TokenStream>();
520
521    let out_rs_path = Path::new(&cargo_out_dir).join("out.rs");
522    let out_tok = quote! {
523      pub mod models {
524        #![allow(unused_imports)]
525        #![allow(clippy::large_enum_variant)]
526
527        use #crate_import::__private::anyhow::{self, anyhow};
528        use #crate_import::__private::serde::{Deserialize, Serialize};
529        use #crate_import::models::chrono;
530
531        #models_out
532      }
533
534      #apis_out
535    };
536    File::create(&out_rs_path)
537      .unwrap_or_else(|err| panic!("failed to create {}: {err}", out_rs_path.to_string_lossy()))
538      .write_all(
539        prettyplease::unparse(
540          &parse2(out_tok.clone())
541            .unwrap_or_else(|err| panic!("failed to parse generated code: {err}\n{out_tok}")),
542        )
543        .as_bytes(),
544      )
545      .unwrap_or_else(|err| {
546        panic!(
547          "failed to write to {}: {err}",
548          out_rs_path.to_string_lossy()
549        )
550      });
551  }
552
553  /// Name of this crate to use for `use` imports.
554  fn crate_use_name(&self) -> Ident {
555    // TODO: support import customization similar to serde's `crate` attribute:
556    // https://serde.rs/container-attrs.html#crate. This also requires a custom model.mustache
557    // since that file embeds the #[serde(crate = "...")] attributes.
558    Ident::new("openapi_lambda", Span::call_site())
559  }
560
561  fn rustfmt(&self, path: &Path) {
562    let rustfmt_result = Command::new("rustfmt")
563      .args(["--edition".as_ref(), "2021".as_ref(), path.as_os_str()])
564      .output()
565      .unwrap_or_else(|err| panic!("failed to run rustfmt: {err}"));
566
567    if !rustfmt_result.status.success() {
568      panic!(
569        "rustfmt failed with status {}:\n{}",
570        rustfmt_result.status,
571        String::from_utf8_lossy(rustfmt_result.stdout.as_slice())
572          + String::from_utf8_lossy(rustfmt_result.stderr.as_slice())
573      );
574    }
575  }
576}
577
578fn description_to_doc_attr<S>(description: &S) -> TokenStream
579where
580  S: AsRef<str>,
581{
582  description
583    .as_ref()
584    .lines()
585    .map(|line| {
586      quote! {
587        #[doc = #line]
588      }
589    })
590    .collect()
591}