Scala avro4s, define SchemaFor for common trait?

49 Views Asked by At

I am trying to define an avro4s schema derivation for a common trait. Example

trait Event
case class UserCreated(name: String, age: Int) extends Event
case class UserDeleted(reason: String) extends Event

Derivation for the concrete classes seems straightforward:

given userCreatedSchemaFor = SchemaFor[UserCreated]
given userDeletedSchemaFor = SchemaFor[UserDeleted]

I would love to achieve something like this fake code:

given eventSchemaFor = new SchemaFor[Event] { 
  CASE (userCreated) => USE userCreatedSchemaFor
  CASE (userDeleted) => USE userDeletedSchemaFor
}

The SchemaFor trait & object expose some apply variants and 'factory' methods but I can't figure out a way to use them to my purpose. Thanks for any help.

2

There are 2 best solutions below

0
ticofab On BEST ANSWER

The answer by @stefanobaghino is correct if you can seal the trait, but in my case I was interested in an unsealed trait. The solution to generate a schema was quite near - avro4s provides a Schema.createUnion which is great for this usecase.

The next problem you might have is encoding for serialization (after all we want to do something useful with this data). For that, the best way I found was the manual creation of an Encoder. Follows a scala-cli script that shows everything put together.

//> using scala 3.3.1
//> using dep com.sksamuel.avro4s::avro4s-core::5.0.9

import org.apache.avro.Schema
import com.sksamuel.avro4s.{AvroOutputStream, AvroSchema, Encoder, SchemaFor}
import java.io.ByteArrayOutputStream
import scala.util.Using

trait Event
case class UserCreated(age: Int) extends Event
case class UserDeleted(reason: String) extends Event
case class Envelope(name: String, event: Event) // <- here I use the trait

// individual schemas
given userCreatedSchema: Schema = AvroSchema[UserCreated] 
given userDeletedSchema: Schema = AvroSchema[UserDeleted]

given unionSchema: Schema = Schema.createUnion(
  AvroSchema[UserCreated],
  AvroSchema[UserDeleted],
)
given unionSchemaFor: SchemaFor[Event] = SchemaFor.from[Event](unionSchema)

val userCreated = UserCreated(18)
val userDeleted = UserDeleted("boring")
val envelopeUserCreated = Envelope("uc", userCreated)
val envelopeUserDeleted = Envelope("ud", userDeleted)

given encoderEvent: Encoder[Event] = (schema: Schema) => {
  case e: UserCreated =>
    summon[Encoder[UserCreated]].encode(userCreatedSchema)(e)
  case e: UserDeleted =>
    summon[Encoder[UserDeleted]].encode(userDeletedSchema)(e)
  case _ => throw new Exception()
}

object Main extends App:
  val byteStream = new ByteArrayOutputStream
  val jsonOutcome = Using(AvroOutputStream.json[Envelope].to(byteStream).build()) { aos =>
    aos.write(envelopeUserCreated)
    aos.write(envelopeUserDeleted)
    aos.flush()
    byteStream.toString
  }
  println(jsonOutcome)
5
stefanobaghino On

If you can seal your trait, sealed trait hierarchies are supported out of the box (see the built-in type mappings in this table) and a few more notes on those here.

The following works as expected without any further intervention:

import com.sksamuel.avro4s.SchemaFor

sealed trait Event
case class UserCreated(name: String, age: Int) extends Event
case class UserDeleted(reason: String) extends Event

val schemaForEvent = SchemaFor[Event]
val schemaForUserCreated = SchemaFor[UserCreated]
val schemaForUserDeleted = SchemaFor[UserDeleted]

You can play around with this code here on Scastie.