main.misk.grpc.reflect.GrpcReflectModule.kt Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of misk-grpc-reflect Show documentation

Open source application container in Kotlin

There is a newer version: 2024.09.17.200749-4708422

package misk.grpc.reflect

import com.google.inject.Provides
import com.squareup.wire.Service
import com.squareup.wire.WireRpc
import com.squareup.wire.reflector.SchemaReflector
import com.squareup.wire.schema.Location
import com.squareup.wire.schema.Schema
import com.squareup.wire.schema.SchemaLoader
import misk.inject.KAbstractModule
import misk.web.WebActionModule
import misk.web.actions.WebAction
import misk.web.actions.WebActionEntry
import okio.FileSystem
import okio.Path.Companion.toPath
import wisp.logging.getLogger
import jakarta.inject.Singleton
import kotlin.reflect.KClass
import kotlin.reflect.full.findAnnotation
import kotlin.reflect.full.functions
import kotlin.reflect.full.isSubclassOf

/**
 * Implements gRPC reflect for all gRPC actions installed in this Misk application.
 *
 * This relies on `.proto` files being included in the `.jar` file. If they're missing, reflection
 * won't work for them.
 */
class GrpcReflectModule : KAbstractModule() {
  override fun configure() {
    install(WebActionModule.create())
  }

  @Provides @Singleton
  fun provideServiceReflector(schema: Schema): SchemaReflector {
    return SchemaReflector(schema)
  }

  /** Interrogate the installed gRPC actions and create a Wire schema from that. */
  @Provides @Singleton
  fun provideSchema(webActions: List): Schema {
    val fileSystem = FileSystem.RESOURCES
    val schemaLoader = SchemaLoader(fileSystem)
    schemaLoader.initRoots(
      sourcePath = toSourceLocations(fileSystem, webActions).toList(),
      protoPath = listOf(Location.get("."))
    )
    schemaLoader.loadExhaustively = true
    return schemaLoader.loadSchema()
  }

  private fun toSourceLocations(
    fileSystem: FileSystem,
    webActions: List
  ): Set {
    val result = mutableSetOf()
    for (webAction in webActions) {
      val wireRpcAnnotation = getWireRpcAnnotation(webAction.actionClass) ?: continue
      val sourceFile = wireRpcAnnotation.sourceFile
      if (sourceFile == "") continue // Generated before @WireRpc.sourceFile existed.
      result += Location.get(".", sourceFile)
    }
    result.removeIf {
      val fileDoesNotExist = !fileSystem.exists(it.path.toPath())
      if (fileDoesNotExist) {
        logger.info("Omitting ${it.path} from ServerReflectionApi; file is not in artifact")
      }
      fileDoesNotExist
    }
    return result
  }

  /** Find `@WireRpc` on a function of one of the supertypes and return it. */
  private fun getWireRpcAnnotation(actionClass: KClass): WireRpc? {
    for (supertype in actionClass.supertypes) {
      val kClass = supertype.classifier as? KClass<*> ?: continue
      if (!kClass.isSubclassOf(Service::class)) continue
      for (function in kClass.functions) {
        val wireRpc = function.findAnnotation()
        if (wireRpc != null) return wireRpc
      }
    }
    return null
  }

  companion object {
    private val logger = getLogger()
  }
}