feat(type-safe-api): generate infinite query hooks for smithy operations with paginated trait

packages/type-safe-api/docs/developer_guides/type-safe-api/typescript_react_query_hooks.md

@@ -155,139 +155,67 @@ export const MyComponent: FC<MyComponentProps> = () => {
 
 ## Paginated Operations
 
-You can generate `useInfiniteQuery` hooks instead of `useQuery` hooks for paginated API operations, by making use of the vendor extension `x-paginated` in your operation in the OpenAPI specification. You must specify both the `inputToken` and `outputToken`, which indicate the properties from the input and output used for pagination. For example in OpenAPI:
-
-```yaml
-paths:
-  /pets:
-    get:
-      x-paginated:
-        # Input property with the token to request the next page
-        inputToken: nextToken
-        # Output property with the token to request the next page
-        outputToken: nextToken
-      parameters:
-        - in: query
-          name: nextToken
-          schema:
-            type: string
-          required: true
-      responses:
-        200:
-          description: Successful response
-          content:
-            application/json:
-              schema:
-                type: object
-                properties:
-                  nextToken:
-                    type: string
-```
-
-In Smithy, until [custom vendor extensions can be rendered via traits](https://github.com/awslabs/smithy/pull/1609), you can add the `x-paginated` vendor extension via `smithyBuildOptions` in your `TypeSafeApiProject`, for example:
-
-=== "TS"
-
-    ```ts
-    new TypeSafeApiProject({
-      model: {
-        language: ModelLanguage.SMITHY,
-        options: {
-          smithy: {
-            serviceName: {
-              namespace: 'com.mycompany',
-              serviceName: 'MyApi',
-            },
-            smithyBuildOptions: {
-              projections: {
-                openapi: {
-                  plugins: {
-                    openapi: {
-                      jsonAdd: {
-                        // Add the x-paginated vendor extension to the GET /pets operation
-                        '/paths/~1pets/get/x-paginated': {
-                          inputToken: 'nextToken',
-                          outputToken: 'nextToken',
-                        },
-                      },
-                    },
-                  },
-                },
-              },
-            },
-          },
-        },
-      },
-      ...
-    });
-    ```
-
-=== "JAVA"
-
-    ```java
-    TypeSafeApiProject.Builder.create()
-            .model(ModelConfiguration.builder()
-                    .language(ModelLanguage.SMITHY)
-                    .options(ModelOptions.builder()
-                            .smithy(SmithyModelOptions.builder()
-                                    .smithyBuildOptions(SmithyBuildOptions.builder()
-                                            .projections(Map.of(
-                                                    "openapi", SmithyProjection.builder()
-                                                            .plugins(Map.of(
-                                                                    "openapi", Map.of(
-                                                                            "jsonAdd", Map.of(
-                                                                                    // Add the x-paginated vendor extension to the GET /pets operation
-                                                                                    "/paths/~1pets/get/x-paginated", Map.of(
-                                                                                            "inputToken", "nextToken",
-                                                                                            "outputToken", "nextToken")
-                                                                            )
-                                                                    )
-                                                            ))
-                                                            .build()
-                                            ))
-                                            .build())
-                                    .build())
-                            .build())
-                    .build())
-            ...
-            .build())
-    .build();
+You can generate `useInfiniteQuery` hooks instead of `useQuery` hooks for paginated API operations by configuring the operation in your model. The configuration requires both an `inputToken` (specifies the input property used to request the next page), and an `outputToken` (specifies the output property in which the pagination token is returned) to be present.
+
+=== "SMITHY"
+
+    In Smithy, annotate your paginated API operations with the `@paginated` trait, making sure both `inputToken` and `outputToken` are specified:
+
+    ```smithy
+    @readonly
+    @http(method: "GET", uri: "/pets")
+    @paginated(inputToken: "nextToken", outputToken: "nextToken", items: "pets") // <- @paginated trait
+    operation ListPets {
+        input := {
+            // Corresponds to inputToken
+            @httpQuery("nextToken")
+            nextToken: String
+        }
+        output := {
+            @required
+            pets: Pets
+
+            // Corresponds to outputToken
+            nextToken: String
+        }
+    }
+
+    list Pets {
+        member: Pet
+    }
     ```
 
-=== "PYTHON"
-
-    ```python
-    TypeSafeApiProject(
-        model={
-            "language": ModelLanguage.SMITHY,
-            "options": {
-                "smithy": {
-                    "service_name": {
-                        "namespace": "com.mycompany",
-                        "service_name": "MyApi"
-                    },
-                    "smithy_build_options": {
-                        "projections": {
-                            "openapi": {
-                                "plugins": {
-                                    "openapi": {
-                                        "json_add": {
-                                            # Add the x-paginated vendor extension to the GET /pets operation
-                                            "/paths/~1pets/get/x-paginated": {
-                                                "input_token": "nextToken",
-                                                "output_token": "nextToken"
-                                            }
-                                        }
-                                    }
-                                }
-                            }
-                        }
-                    }
-                }
-            }
-        },
-        ...
-    )
+=== "OPENAPI"
+
+    In OpenAPI, use the `x-paginaged` vendor extension in your operation, making sure both `inputToken` and `outputToken` are specified:
+
+    ```yaml
+    paths:
+      /pets:
+        get:
+          x-paginated:
+            # Input property with the token to request the next page
+            inputToken: nextToken
+            # Output property with the token to request the next page
+            outputToken: nextToken
+          parameters:
+            - in: query
+              name: nextToken
+              schema:
+                type: string
+              required: true
+          responses:
+            200:
+              description: Successful response
+              content:
+                application/json:
+                  schema:
+                    type: object
+                    properties:
+                      nextToken:
+                        type: string
+                      pets:
+                        $ref: "#/components/schemas/Pets"
     ```
 
 ## Custom QueryClient

packages/type-safe-api/scripts/generators/generate

@@ -12,6 +12,7 @@ openapi_normalizer=''
 src_dir='src'
 extra_vendor_extensions=''
 generate_alias_as_model=''
+smithy_json_path=''
 while [[ "$#" -gt 0 ]]; do case $1 in
   --spec-path) spec_path="$2"; shift;;
   --output-path) output_path="$2"; shift;;
@@ -21,6 +22,7 @@ while [[ "$#" -gt 0 ]]; do case $1 in
   --openapi-normalizer) openapi_normalizer="$2"; shift;;
   --src-dir) src_dir="$2"; shift;;
   --extra-vendor-extensions) extra_vendor_extensions="$2"; shift;;
+  --smithy-json-path) smithy_json_path="$2"; shift;;
   --generate-alias-as-model) generate_alias_as_model='true'; ;;
 esac; shift; done
 
@@ -53,11 +55,11 @@ install_packages \
   [email protected]
 
 # Pre-process the spec to add any extra vendor extensions used for code generation only.
-# This allows for custom parameters to be passed to the code generation templates.
+# This allows for custom parameters to be passed to the code generation templates, and smithy traits to be accessed during codegen
 log "preprocess spec :: $spec_path"
 processed_spec_path="$tmp_dir/.preprocessed-api.json"
 cp $script_dir/pre-process-spec.ts .
-run_command ts-node pre-process-spec.ts --specPath="$working_dir/$spec_path" --outputSpecPath="$processed_spec_path" --extraVendorExtensions="$extra_vendor_extensions"
+run_command ts-node pre-process-spec.ts --specPath="$working_dir/$spec_path" --outputSpecPath="$processed_spec_path" --extraVendorExtensions="$extra_vendor_extensions" ${smithy_json_path:+"--smithyJsonPath=$working_dir/$smithy_json_path"}
 
 # Support a special placeholder of {{src}} in config.yaml to ensure our custom templates get written to the correct folder
 sed 's|{{src}}|'"$src_dir"'|g' config.yaml > config.final.yaml

packages/type-safe-api/scripts/generators/pre-process-spec.ts

@@ -5,11 +5,23 @@ import * as path from "path";
 import SwaggerParser from "@apidevtools/swagger-parser";
 import { parse } from "ts-command-line-args";
 
+// Smithy HTTP trait is used to map Smithy operations to their location in the spec
+const SMITHY_HTTP_TRAIT_ID = "smithy.api#http";
+
+// Maps traits to specific vendor extensions which we also support specifying in OpenAPI
+const TRAIT_TO_SUPPORTED_OPENAPI_VENDOR_EXTENSION: { [key: string]: string } = {
+  "smithy.api#paginated": "x-paginated",
+};
+
 interface Arguments {
   /**
    * Path to the spec to preprocess
    */
   readonly specPath: string;
+  /**
+   * Path to the smithy json model
+   */
+  readonly smithyJsonPath?: string;
   /**
    * JSON string containing extra vendor extensions to add to the spec
    */
@@ -20,14 +32,58 @@ interface Arguments {
   readonly outputSpecPath: string;
 }
 
+interface SmithyOperationDetails {
+  readonly id: string;
+  readonly method: string;
+  readonly path: string;
+  readonly traits: { [key: string]: any };
+}
+
 void (async () => {
   const args = parse<Arguments>({
     specPath: { type: String },
+    smithyJsonPath: { type: String, optional: true },
     extraVendorExtensions: { type: String, optional: true },
     outputSpecPath: { type: String },
   });
 
-  const spec = await SwaggerParser.bundle(args.specPath);
+  const spec = (await SwaggerParser.bundle(args.specPath)) as any;
+
+  if (args.smithyJsonPath) {
+    // Read the operations out of the Smithy model
+    const smithyModel = JSON.parse(
+      fs.readFileSync(args.smithyJsonPath, "utf-8")
+    );
+    const operations: SmithyOperationDetails[] = Object.entries(
+      smithyModel.shapes
+    )
+      .filter(
+        ([, shape]: [string, any]) =>
+          shape.type === "operation" &&
+          shape.traits &&
+          SMITHY_HTTP_TRAIT_ID in shape.traits
+      )
+      .map(([id, shape]: [string, any]) => ({
+        id,
+        method: shape.traits[SMITHY_HTTP_TRAIT_ID].method?.toLowerCase(),
+        path: shape.traits[SMITHY_HTTP_TRAIT_ID].uri,
+        traits: shape.traits,
+      }));
+
+    // Apply all operation-level traits as vendor extensions to the relevant operation in the spec
+    operations.forEach((operation) => {
+      if (spec.paths?.[operation.path]?.[operation.method]) {
+        Object.entries(operation.traits).forEach(([traitId, value]) => {
+          // By default, we use x-<fully_qualified_trait_id> for the vendor extension, but for extensions we support
+          // directly from OpenAPI we apply a mapping (rather than repeat ourselves in the mustache templates).
+          const vendorExtension =
+            TRAIT_TO_SUPPORTED_OPENAPI_VENDOR_EXTENSION[traitId] ??
+            `x-${traitId}`;
+          spec.paths[operation.path][operation.method][vendorExtension] = value;
+        });
+      }
+    });
+  }
 
   const processedSpec = {
     ...spec,

packages/type-safe-api/src/project/codegen/components/utils.ts

@@ -79,6 +79,10 @@ export interface GenerationOptions {
    * @default true
    */
   readonly generateAliasAsModel?: boolean;
+  /**
+   * The path to the json smithy model file, if available
+   */
+  readonly smithyJsonPath?: string;
 }
 
 /**
@@ -131,10 +135,14 @@ export const buildInvokeOpenApiGeneratorCommandArgs = (
   const generateAliasAsModel =
     options.generateAliasAsModel ?? true ? " --generate-alias-as-model" : "";
 
+  const smithyJsonPath = options.smithyJsonPath
+    ? ` --smithy-json-path ${options.smithyJsonPath}`
+    : "";
+
   const specPath = options.specPath;
   const outputPath = ".";
 
-  return `--generator ${options.generator} --spec-path ${specPath} --output-path ${outputPath} --generator-dir ${options.generatorDirectory} --src-dir ${srcDir}${additionalProperties}${normalizers}${extensions}${generateAliasAsModel}`;
+  return `--generator ${options.generator} --spec-path ${specPath} --output-path ${outputPath} --generator-dir ${options.generatorDirectory} --src-dir ${srcDir}${smithyJsonPath}${additionalProperties}${normalizers}${extensions}${generateAliasAsModel}`;
 };
 
 /**

packages/type-safe-api/src/project/codegen/generate.ts

@@ -60,6 +60,10 @@ export interface GenerateProjectsOptions {
    * We use the parsed spec such that refs are resolved to support multi-file specs
    */
   readonly parsedSpecPath: string;
+  /**
+   * Path to the Smithy json model, if applicable
+   */
+  readonly smithyJsonModelPath?: string;
 }
 
 /**
@@ -140,6 +144,7 @@ const generateRuntimeProject = (
     outdir: path.join(options.generatedCodeDir, language),
     specPath: options.parsedSpecPath,
     parent: options.parent,
+    smithyJsonModelPath: options.smithyJsonModelPath,
   };
 
   switch (language) {
@@ -213,6 +218,7 @@ export const generateInfraProject = (
     outdir: path.join(options.generatedCodeDir, language),
     specPath: options.parsedSpecPath,
     parent: options.parent,
+    smithyJsonModelPath: options.smithyJsonModelPath,
   };
 
   switch (language) {
@@ -314,6 +320,7 @@ const generateLibraryProject = (
     outdir: path.join(options.generatedCodeDir, library),
     specPath: options.parsedSpecPath,
     parent: options.parent,
+    smithyJsonModelPath: options.smithyJsonModelPath,
   };
 
   switch (library) {