support Avro fixed type

spotify · Feb 25, 2021 · bec818d · bec818d
1 parent ca3927a
commit bec818d
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 2 deletions.
diff --git a/avro/src/main/scala/magnolify/avro/AvroType.scala b/avro/src/main/scala/magnolify/avro/AvroType.scala
@@ -32,6 +32,7 @@ import scala.annotation.{implicitNotFound, StaticAnnotation}
 import scala.collection.concurrent
 import scala.language.experimental.macros
 import scala.language.implicitConversions
+import scala.reflect.ClassTag
 
 class doc(doc: String) extends StaticAnnotation with Serializable {
   override def toString: String = doc
@@ -306,4 +307,21 @@ object AvroField {
     logicalType[String](LogicalTypes.uuid())(ju.UUID.fromString)(_.toString)
   implicit val afDate: AvroField[LocalDate] =
     logicalType[Int](LogicalTypes.date())(LocalDate.ofEpochDay(_))(_.toEpochDay.toInt)
+
+  def fixed[T: ClassTag](
+    size: Int
+  )(f: Array[Byte] => T)(g: T => Array[Byte])(implicit an: AnnotationType[T]): AvroField[T] =
+    new AvroField[T] {
+      override type FromT = GenericFixed
+      override type ToT = GenericFixed
+
+      override protected def schemaString(cm: CaseMapper): String = {
+        val n = ReflectionUtils.name[T]
+        val ns = ReflectionUtils.namespace[T]
+        Schema.createFixed(n, getDoc(an.annotations, n), ns, size).toString
+      }
+
+      override def from(v: GenericFixed)(cm: CaseMapper): T = f(v.bytes())
+      override def to(v: T)(cm: CaseMapper): GenericFixed = new GenericData.Fixed(schema(cm), g(v))
+    }
 }
diff --git a/avro/src/test/scala/magnolify/avro/test/AvroTypeSuite.scala b/avro/src/test/scala/magnolify/avro/test/AvroTypeSuite.scala
@@ -133,6 +133,23 @@ class AvroTypeSuite extends MagnolifySuite {
     }
   }
 
+  {
+    implicit val arbCountryCode: Arbitrary[CountryCode] = Arbitrary(
+      Gen.oneOf("US", "UK", "CA", "MX").map(CountryCode(_))
+    )
+    implicit val afCountryCode: AvroField[CountryCode] =
+      AvroField.fixed[CountryCode](2)(bs => CountryCode(new String(bs)))(cc => cc.code.getBytes)
+    test[Fixed]
+
+    test("FixedDoc") {
+      val at = ensureSerializable(AvroType[Fixed])
+      val schema = at.schema.getField("countryCode").schema
+      assertEquals(schema.getName, "CountryCode")
+      assertEquals(schema.getNamespace, "magnolify.avro.test")
+      assertEquals(schema.getDoc, "Fixed with doc")
+    }
+  }
+
   test("AvroDoc") {
     val at = ensureSerializable(AvroType[AvroDoc])
     val schema = at.schema
@@ -268,6 +285,10 @@ case class LogicalMillis(bd: BigDecimal, i: Instant, t: LocalTime, dt: LocalDate
 case class LogicalBigQuery(bd: BigDecimal, i: Instant, t: LocalTime, dt: LocalDateTime)
 case class BigDec(bd: BigDecimal)
 
+@doc("Fixed with doc")
+case class CountryCode(code: String)
+case class Fixed(countryCode: CountryCode)
+
 @doc("Avro with doc")
 case class AvroDoc(@doc("string") s: String, @doc("integers") i: Integers)
 

diff --git a/docs/avro.md b/docs/avro.md
@@ -5,16 +5,21 @@ AvroType
 
 ```scala
 import java.net.URI
-case class Inner(long: Long, str: String, uri: URI)
+case class CountryCode(code: String)
+case class Inner(long: Long, str: String, uri: URI, cc: CountryCode)
 case class Outer(inner: Inner)
-val record = Outer(Inner(1L, "hello", URI.create("https://www.spotify.com")))
+val record = Outer(Inner(1L, "hello", URI.create("https://www.spotify.com"), "US"))
 
 import magnolify.avro._
 import org.apache.avro.generic.GenericRecord
 
 // Encode custom type URI as String
 implicit val uriField = AvroField.from[String](URI.create)(_.toString)
 
+// Encode country code as fixed type
+implicit val afCountryCode =
+  AvroField.fixed[CountryCode](2)(bs => CountryCode(new String(bs)))(cc => cc.code.getBytes)
+
 val avroType = AvroType[Outer]
 val genericRecord: GenericRecord = avroType.to(record)
 val copy: Outer = avroType.from(genericRecord)