Respect case mapper when generating avro schema (#787)

spotify · Sep 12, 2023 · 35de174 · 35de174
1 parent 808f103
commit 35de174
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 24 deletions.
diff --git a/parquet/src/main/scala/magnolify/parquet/ParquetField.scala b/parquet/src/main/scala/magnolify/parquet/ParquetField.scala
@@ -43,8 +43,9 @@ sealed trait ParquetField[T] extends Serializable {
     schemaCache.getOrElseUpdate(cm.uuid, buildSchema(cm))
 
   val hasAvroArray: Boolean = false
-  val fieldDocs: Map[String, String]
-  val typeDoc: Option[String]
+  def fieldDocs(cm: CaseMapper): Map[String, String]
+  def typeDoc: Option[String]
+
   protected val isGroup: Boolean = false
   protected def isEmpty(v: T): Boolean
   def write(c: RecordConsumer, v: T)(cm: CaseMapper): Unit
@@ -87,7 +88,7 @@ object ParquetField {
             override def get: T = inner.get(b => caseClass.construct(_ => b.head))
           }
         }
-        override val fieldDocs: Map[String, String] = Map.empty
+        override def fieldDocs(cm: CaseMapper): Map[String, String] = Map.empty
         override val typeDoc: Option[String] = None
       }
     } else {
@@ -101,10 +102,10 @@ object ParquetField {
 
         override val hasAvroArray: Boolean = caseClass.parameters.exists(_.typeclass.hasAvroArray)
 
-        override val fieldDocs: Map[String, String] =
+        override def fieldDocs(cm: CaseMapper): Map[String, String] =
           caseClass.parameters.flatMap { param =>
-            val label = param.label
-            val nestedDocs = param.typeclass.fieldDocs.map { case (k, v) =>
+            val label = cm.map(param.label)
+            val nestedDocs = param.typeclass.fieldDocs(cm).map { case (k, v) =>
               s"$label.$k" -> v
             }
 
@@ -198,7 +199,7 @@ object ParquetField {
 
   sealed trait Primitive[T] extends ParquetField[T] {
     override protected def isEmpty(v: T): Boolean = false
-    override val fieldDocs: Map[String, String] = Map.empty
+    override def fieldDocs(cm: CaseMapper): Map[String, String] = Map.empty
     override val typeDoc: Option[String] = None
     type ParquetT <: Comparable[ParquetT]
   }
@@ -281,7 +282,7 @@ object ParquetField {
         Schema.setRepetition(t.schema(cm), Repetition.OPTIONAL)
       override protected def isEmpty(v: Option[T]): Boolean = v.isEmpty
 
-      override val fieldDocs: Map[String, String] = t.fieldDocs
+      override def fieldDocs(cm: CaseMapper): Map[String, String] = t.fieldDocs(cm)
 
       override val typeDoc: Option[String] = None
 
@@ -359,7 +360,7 @@ object ParquetField {
         }
       }
 
-      override val fieldDocs: Map[String, String] = t.fieldDocs
+      override def fieldDocs(cm: CaseMapper): Map[String, String] = t.fieldDocs(cm)
 
       override val typeDoc: Option[String] = None
     }

diff --git a/parquet/src/main/scala/magnolify/parquet/ParquetType.scala b/parquet/src/main/scala/magnolify/parquet/ParquetType.scala
@@ -90,7 +90,8 @@ object ParquetType {
         @transient override lazy val avroSchema: AvroSchema = {
           val s = new AvroSchemaConverter().convert(schema)
           // add doc to avro schema
-          SchemaUtil.deepCopy(s, f.typeDoc, f.fieldDocs.get)
+          val fieldDocs = f.fieldDocs(cm)
+          SchemaUtil.deepCopy(s, f.typeDoc, fieldDocs.get)
         }
 
         override val avroCompat: Boolean =

diff --git a/parquet/src/test/scala/magnolify/parquet/AvroParquetSuite.scala b/parquet/src/test/scala/magnolify/parquet/AvroParquetSuite.scala
@@ -21,7 +21,7 @@ import cats._
 import magnolify.cats.auto._
 import magnolify.cats.TestEq._
 import magnolify.avro.AvroType
-import magnolify.shared.doc
+import magnolify.shared.{doc, CaseMapper}
 import magnolify.avro.unsafe._
 import magnolify.parquet.unsafe._
 import magnolify.parquet.ParquetArray.AvroCompat._
@@ -44,8 +44,24 @@ import org.scalacheck._
 import scala.reflect.ClassTag
 
 class AvroParquetSuite extends MagnolifySuite {
+
+  private def test[T: Arbitrary: ClassTag]()(implicit
+    at: AvroType[T],
+    tpe: ParquetType[T],
+    eq: Eq[T]
+  ): Unit = test[T](List.empty)
+
   private def test[T: Arbitrary: ClassTag](
-    schemaErrors: List[String] = List.empty
+    schemaErrors: List[String]
+  )(implicit
+    at: AvroType[T],
+    tpe: ParquetType[T],
+    eq: Eq[T]
+  ): Unit = test[T](className[T], schemaErrors)
+
+  private def test[T: Arbitrary](
+    name: String,
+    schemaErrors: List[String]
   )(implicit
     at: AvroType[T],
     tpe: ParquetType[T],
@@ -56,8 +72,6 @@ class AvroParquetSuite extends MagnolifySuite {
     val avroSchema = tpe.avroSchema
     val pt = ensureSerializable(tpe)
 
-    val name = className[T]
-
     property(s"$name.avro2parquet") {
       Prop.forAll { t: T =>
         val r = at(t)
@@ -192,13 +206,29 @@ class AvroParquetSuite extends MagnolifySuite {
       "root.nested 'namespace' are different 'magnolify.parquet' != 'null'"
     )
   )
+
+  {
+    val upperCaseMapper = CaseMapper(_.toUpperCase())
+    implicit val at = AvroType[AvroParquetWithNestedAnnotations](upperCaseMapper)
+    implicit val pt = ParquetType[AvroParquetWithNestedAnnotations](upperCaseMapper)
+    test[AvroParquetWithNestedAnnotations](
+      "AvroParquetWithNestedAnnotations and CaseMapper",
+      List(
+        "root.NESTED 'name' are different 'AvroParquetWithAnnotations' != 'NESTED'",
+        "root.NESTED 'doc' are different 'Should be ignored' != 'null'",
+        "root.NESTED 'namespace' are different 'magnolify.parquet' != 'null'"
+      )
+    )
+  }
+
   test[AvroParquetWithAnnotationsAndOptions](
     List(
       "root.nested 'name' are different 'AvroParquetWithAnnotations' != 'nested'",
       "root.nested 'doc' are different 'Should be ignored' != 'null'",
       "root.nested 'namespace' are different 'magnolify.parquet' != 'null'"
     )
   )
+
   test[AvroParquetWithAnnotationsAndLists](
     List(
       "root.nested 'name' are different 'AvroParquetWithAnnotations' != 'array'",

diff --git a/parquet/src/test/scala/magnolify/parquet/ParquetTypeSuite.scala b/parquet/src/test/scala/magnolify/parquet/ParquetTypeSuite.scala
@@ -102,23 +102,31 @@ class ParquetTypeSuite extends MagnolifySuite {
   test("ParquetDoc") {
     ensureSerializable(ParquetType[ParquetNestedDoc])
     val pf = ParquetField[ParquetNestedDoc]
-
-    assert(pf.fieldDocs("pd") == "nested")
-    assert(pf.fieldDocs("pd.i") == "integers")
-    assert(pf.fieldDocs("pd.s") == "string")
-    assert(pf.fieldDocs("i") == "integers")
     assert(pf.typeDoc.contains("Parquet with doc"))
+
+    val fieldDocs = pf.fieldDocs(CaseMapper.identity)
+    assert(fieldDocs("pd") == "nested")
+    assert(fieldDocs("pd.i") == "integers")
+    assert(fieldDocs("pd.s") == "string")
+    assert(fieldDocs("i") == "integers")
+
+    val upperFieldDocs = pf.fieldDocs(CaseMapper(_.toUpperCase()))
+    assert(upperFieldDocs("PD") == "nested")
+    assert(upperFieldDocs("PD.I") == "integers")
+    assert(upperFieldDocs("PD.S") == "string")
+    assert(upperFieldDocs("I") == "integers")
   }
 
   test("ParquetDocWithNestedList") {
     ensureSerializable(ParquetType[ParquetNestedListDoc])
     val pf = ParquetField[ParquetNestedListDoc]
-
-    assert(pf.fieldDocs("pd") == "nested")
-    assert(pf.fieldDocs("pd.i") == "integers")
-    assert(pf.fieldDocs("pd.s") == "string")
-    assert(pf.fieldDocs("i") == "integers")
     assert(pf.typeDoc.contains("Parquet with doc with nested list"))
+
+    val fieldDocs = pf.fieldDocs(CaseMapper.identity)
+    assert(fieldDocs("pd") == "nested")
+    assert(fieldDocs("pd.i") == "integers")
+    assert(fieldDocs("pd.s") == "string")
+    assert(fieldDocs("i") == "integers")
   }
 
   // Precision = number of digits, so 5 means -99999 to 99999