Smithy4s Preprocessor Split Services

SplitUp.scala

package preprocessors

import software.amazon.smithy.build._
import software.amazon.smithy.model._
import software.amazon.smithy.model.shapes._
import software.amazon.smithy.model.traits._
import scala.collection.JavaConverters._

final class SplitUp extends ProjectionTransformer {

  // Add new tags and services here, nothing else in the
  // file should (probably) need to be modified
  // TagName -> ServiceName
  private val serviceGroupings = Map(
    "auth" -> "TestAuthService"
  )

  def getName() = {
    "SplitUp"
  }

  def transform(ctx: TransformContext): Model = {
    val model = ctx.getModel()
    val startingService = model
      .expectShape(ShapeId.fromParts("test.api", "TestService"))
      .asServiceShape
      .get
    val startingOperations = startingService
      .getOperations()
      .asScala
      .map(model.expectShape(_).asOperationShape.get)

    val getOpsForTag: Option[String] => Set[OperationShape] = maybeTag =>
      startingOperations.filter { op =>
        val maybeTags = op.getTrait(classOf[TagsTrait])
        maybeTag match {
          case None =>
            !maybeTags
              .isPresent() // let through if maybeTag is NOT defined and NO tag trait is defined for the op
          case Some(tag) =>
            if (maybeTags.isPresent) maybeTags.get.getTags().contains(tag)
            else false
        }
      }.toSet

    val leftOvers = getOpsForTag(None)

    val updatedMainService =
      startingService
        .toBuilder()
        .operations(leftOvers.map(_.toShapeId).asJava)
        .build

    val newServices = serviceGroupings.toList.map {
      case (tagName, serviceName) =>
        val ops = getOpsForTag(Some(tagName))

        ServiceShape
          .builder()
          .id(ShapeId.fromParts("test.api", serviceName))
          .operations(ops.map(_.toShapeId).asJava)
          .errors(startingService.getErrors())
          .addTrait(new alloy.SimpleRestJsonTrait())
          .build
    }

    val allNewServices = newServices :+ updatedMainService

    val updatedModel =
      model.toBuilder().addShapes(allNewServices: _*).build
    updatedModel
  }
}