@kubb/plugin-oas/src/OperationGenerator.ts

OpenAPI Specification (OAS) plugin for Kubb, providing core functionality for parsing and processing OpenAPI/Swagger schemas for code generation.

import { pascalCase } from '@internals/utils'
import type { AsyncEventEmitter, FileMetaBase, KubbEvents, Plugin, PluginFactoryOptions, PluginManager } from '@kubb/core'
import type { KubbFile } from '@kubb/fabric-core/types'
import type { contentType, HttpMethod, Oas, OasTypes, Operation, SchemaObject } from '@kubb/oas'
import type { Fabric } from '@kubb/react-fabric/types'
import pLimit from 'p-limit'
import type { CoreGenerator } from './generators/createGenerator.ts'
import type { ReactGenerator } from './generators/createReactGenerator.ts'
import type { Generator, Version } from './generators/types.ts'
import type { Exclude, Include, OperationSchemas, Override } from './types.ts'
import { withRequiredRequestBodySchema } from './utils/requestBody.ts'
import { buildOperation, buildOperations } from './utils.tsx'

export type OperationMethodResult<TFileMeta extends FileMetaBase> = Promise<KubbFile.File<TFileMeta> | Array<KubbFile.File<TFileMeta>> | null>

type Context<TOptions, TPluginOptions extends PluginFactoryOptions> = {
  fabric: Fabric
  oas: Oas
  exclude: Array<Exclude> | undefined
  include: Array<Include> | undefined
  override: Array<Override<TOptions>> | undefined
  contentType: contentType | undefined
  pluginManager: PluginManager
  events?: AsyncEventEmitter<KubbEvents>
  /**
   * Current plugin
   */
  plugin: Plugin<TPluginOptions>
  mode: KubbFile.Mode
  UNSTABLE_NAMING?: true
}

export class OperationGenerator<TPluginOptions extends PluginFactoryOptions = PluginFactoryOptions, TFileMeta extends FileMetaBase = FileMetaBase> {
  #options: TPluginOptions['resolvedOptions']
  #context: Context<TPluginOptions['resolvedOptions'], TPluginOptions>

  constructor(options: TPluginOptions['resolvedOptions'], context: Context<TPluginOptions['resolvedOptions'], TPluginOptions>) {
    this.#options = options
    this.#context = context
  }

  get options(): TPluginOptions['resolvedOptions'] {
    return this.#options
  }

  set options(options: TPluginOptions['resolvedOptions']) {
    this.#options = { ...this.#options, ...options }
  }

  get context(): Context<TPluginOptions['resolvedOptions'], TPluginOptions> {
    return this.#context
  }
  #matchesPattern(operation: Operation, method: HttpMethod, type: string, pattern: RegExp | string): boolean {
    switch (type) {
      case 'tag':
        return operation.getTags().some((tag) => tag.name.match(pattern))
      case 'operationId':
        return !!operation.getOperationId({ friendlyCase: true }).match(pattern)
      case 'path':
        return !!operation.path.match(pattern)
      case 'method':
        return !!method.match(pattern)
      case 'contentType':
        return !!operation.getContentType().match(pattern)
      default:
        return false
    }
  }

  getOptions(operation: Operation, method: HttpMethod): Partial<TPluginOptions['resolvedOptions']> {
    const { override = [] } = this.context

    return override.find(({ pattern, type }) => this.#matchesPattern(operation, method, type, pattern))?.options || {}
  }

  #isExcluded(operation: Operation, method: HttpMethod): boolean {
    const { exclude = [] } = this.context

    return exclude.some(({ pattern, type }) => this.#matchesPattern(operation, method, type, pattern))
  }

  #isIncluded(operation: Operation, method: HttpMethod): boolean {
    const { include = [] } = this.context

    return include.some(({ pattern, type }) => this.#matchesPattern(operation, method, type, pattern))
  }

  getSchemas(
    operation: Operation,
    {
      resolveName = (name) => name,
    }: {
      resolveName?: (name: string) => string
    } = {},
  ): OperationSchemas {
    const operationId = operation.getOperationId({ friendlyCase: true })
    const operationName = pascalCase(operationId)

    const resolveKeys = (schema?: SchemaObject) => (schema?.properties ? Object.keys(schema.properties) : undefined)

    const pathParamsSchema = this.context.oas.getParametersSchema(operation, 'path')
    const queryParamsSchema = this.context.oas.getParametersSchema(operation, 'query')
    const headerParamsSchema = this.context.oas.getParametersSchema(operation, 'header')
    const requestSchema = this.context.oas.getRequestSchema(operation)
    const statusCodes = operation.getResponseStatusCodes().map((statusCode) => {
      const name = statusCode === 'default' ? 'error' : statusCode
      const schema = this.context.oas.getResponseSchema(operation, statusCode)
      const keys = resolveKeys(schema)

      return {
        name: this.context.UNSTABLE_NAMING ? resolveName(pascalCase(`${operationId} status ${name}`)) : resolveName(pascalCase(`${operationId} ${name}`)),
        description: (operation.getResponseByStatusCode(statusCode) as OasTypes.ResponseObject)?.description,
        schema,
        operation,
        operationName,
        statusCode: name === 'error' ? undefined : Number(statusCode),
        keys,
        keysToOmit: keys?.filter((key) => (schema?.properties?.[key] as OasTypes.SchemaObject)?.writeOnly),
      }
    })

    const successful = statusCodes.filter((item) => item.statusCode?.toString().startsWith('2'))
    const errors = statusCodes.filter((item) => item.statusCode?.toString().startsWith('4') || item.statusCode?.toString().startsWith('5'))

    const request = withRequiredRequestBodySchema(
      requestSchema
        ? {
            name: this.context.UNSTABLE_NAMING
              ? resolveName(pascalCase(`${operationId} RequestData`))
              : resolveName(pascalCase(`${operationId} ${operation.method === 'get' ? 'queryRequest' : 'mutationRequest'}`)),
            description: (operation.schema.requestBody as OasTypes.RequestBodyObject)?.description,
            operation,
            operationName,
            schema: requestSchema,
            keys: resolveKeys(requestSchema),
            keysToOmit: resolveKeys(requestSchema)?.filter((key) => (requestSchema.properties?.[key] as OasTypes.SchemaObject)?.readOnly),
          }
        : undefined,
    )

    return {
      pathParams: pathParamsSchema
        ? {
            name: resolveName(pascalCase(`${operationId} PathParams`)),
            operation,
            operationName,
            schema: pathParamsSchema,
            keys: resolveKeys(pathParamsSchema),
          }
        : undefined,
      queryParams: queryParamsSchema
        ? {
            name: resolveName(pascalCase(`${operationId} QueryParams`)),
            operation,
            operationName,
            schema: queryParamsSchema,
            keys: resolveKeys(queryParamsSchema) || [],
          }
        : undefined,
      headerParams: headerParamsSchema
        ? {
            name: resolveName(pascalCase(`${operationId} HeaderParams`)),
            operation,
            operationName,
            schema: headerParamsSchema,
            keys: resolveKeys(headerParamsSchema),
          }
        : undefined,
      request,
      response: {
        name: this.context.UNSTABLE_NAMING
          ? resolveName(pascalCase(`${operationId} ResponseData`))
          : resolveName(pascalCase(`${operationId} ${operation.method === 'get' ? 'queryResponse' : 'mutationResponse'}`)),
        operation,
        operationName,
        schema: {
          oneOf: successful.map((item) => ({ ...item.schema, $ref: item.name })) || undefined,
        } as SchemaObject,
      },
      responses: successful,
      errors,
      statusCodes,
    }
  }

  async getOperations(): Promise<Array<{ path: string; method: HttpMethod; operation: Operation }>> {
    const { oas } = this.context

    const paths = oas.getPaths()

    return Object.entries(paths).flatMap(([path, methods]) =>
      Object.entries(methods)
        .map((values) => {
          const [method, operation] = values as [HttpMethod, Operation]
          if (this.#isExcluded(operation, method)) {
            return null
          }

          if (this.context.include && !this.#isIncluded(operation, method)) {
            return null
          }

          return operation ? { path, method: method as HttpMethod, operation } : null
        })
        .filter((x): x is { path: string; method: HttpMethod; operation: Operation } => x !== null),
    )
  }

  async build(...generators: Array<Generator<TPluginOptions, Version>>): Promise<Array<KubbFile.File<TFileMeta>>> {
    const operations = await this.getOperations()

    // Increased parallelism for better performance
    // - generatorLimit increased from 1 to 3 to allow parallel generator processing
    // - operationLimit increased from 10 to 30 to process more operations concurrently
    const generatorLimit = pLimit(3)
    const operationLimit = pLimit(30)

    this.context.events?.emit('debug', {
      date: new Date(),
      logs: [`Building ${operations.length} operations`, `  • Generators: ${generators.length}`],
    })

    const writeTasks = generators.map((generator) =>
      generatorLimit(async () => {
        if (generator.version === '2') {
          return []
        }

        // After the v2 guard above, all generators here are v1
        const v1Generator = generator as ReactGenerator<TPluginOptions, '1'> | CoreGenerator<TPluginOptions, '1'>

        const operationTasks = operations.map(({ operation, method }) =>
          operationLimit(async () => {
            const options = this.getOptions(operation, method)

            if (v1Generator.type === 'react') {
              await buildOperation(operation, {
                config: this.context.pluginManager.config,
                fabric: this.context.fabric,
                Component: v1Generator.Operation,
                generator: this,
                plugin: {
                  ...this.context.plugin,
                  options: {
                    ...this.options,
                    ...options,
                  },
                },
              })

              return []
            }

            const result = await v1Generator.operation?.({
              generator: this,
              config: this.context.pluginManager.config,
              operation,
              plugin: {
                ...this.context.plugin,
                options: {
                  ...this.options,
                  ...options,
                },
              },
            })

            return result ?? []
          }),
        )

        const operationResults = await Promise.all(operationTasks)
        const opResultsFlat = operationResults.flat()

        if (v1Generator.type === 'react') {
          await buildOperations(
            operations.map((op) => op.operation),
            {
              fabric: this.context.fabric,
              config: this.context.pluginManager.config,
              Component: v1Generator.Operations,
              generator: this,
              plugin: this.context.plugin,
            },
          )

          return []
        }

        const operationsResult = await v1Generator.operations?.({
          generator: this,
          config: this.context.pluginManager.config,
          operations: operations.map((op) => op.operation),
          plugin: this.context.plugin,
        })

        return [...opResultsFlat, ...(operationsResult ?? [])] as unknown as KubbFile.File<TFileMeta>
      }),
    )

    const nestedResults = await Promise.all(writeTasks)

    return nestedResults.flat()
  }
}