Fix bugs

spotify · Feb 26, 2024 · 9ecd001 · 9ecd001
1 parent 30cdf43
commit 9ecd001
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 32 deletions.
diff --git a/src/main/scala/com/spotify/scio/AnnotationTypeInjector.scala b/src/main/scala/com/spotify/scio/AnnotationTypeInjector.scala
@@ -19,7 +19,7 @@ package com.spotify.scio
 
 import java.io.File
 import java.nio.charset.Charset
-import java.nio.file.{Paths, Files as JFiles}
+import java.nio.file.{Files as JFiles, Paths}
 import com.google.common.base.Charsets
 import com.google.common.hash.Hashing
 import com.google.common.io.Files
@@ -34,6 +34,15 @@ import scala.collection.mutable
 
 object AnnotationTypeInjector {
   private val Log = Logger.getInstance(classOf[AnnotationTypeInjector])
+
+  // case classes implement Product trait
+  val CaseClassSuper: String = "_root_.scala.Product"
+  val CaseClassFunctions: Seq[String] = Seq(
+    "def productArity: _root_.scala.Int = ???",
+    "def productElement(n: _root_.scala.Int): _root_.scala.Any = ???",
+    "def canEqual(x: _root_.scala.Any): _root_.scala.Boolean = ???"
+  )
+
   private val CaseClassArgs = """case\s+class\s+[^(]+\((.*)\).*""".r
   private val TypeArg = """[a-zA-Z0-9_$]+\s*:\s*[a-zA-Z0-9._$]+([\[(](.*?)[)\]]+)?""".r
   private val AlertEveryMissedXInvocations = 5

diff --git a/src/main/scala/com/spotify/scio/AvroTypeInjector.scala b/src/main/scala/com/spotify/scio/AvroTypeInjector.scala
@@ -26,8 +26,7 @@ object AvroTypeInjector {
     s"$AvroTNamespace.fromPath",
     s"$AvroTNamespace.toSchema"
   )
-  private val CaseClassSuper =
-    "_root_.scala.Product"
+
   private val HasAvroAnnotationSuper =
     "_root_.com.spotify.scio.avro.types.AvroType.HasAvroAnnotation"
 
@@ -46,19 +45,14 @@ final class AvroTypeInjector extends AnnotationTypeInjector {
   override def injectFunctions(source: ScTypeDefinition): Seq[String] =
     source match {
       case c: ScClass if avroAnnotation(c).isDefined =>
-        val result = for {
-          cc <- Option(c.containingClass)
-          qn <- Option(cc.getQualifiedName)
+        val fields = for {
+          cc <- Option(c.containingClass).toSeq
+          qn <- Option(cc.getQualifiedName).toSeq
           parent = qn.init
-          defs <- {
-            generatedCaseClasses(parent, c)
-              .find(_.contains(HasAvroAnnotationSuper))
-              .map(getApplyPropsSignature)
-              .map(v => s"def $v = ???")
-          }
-        } yield defs
-
-        result.toSeq
+          cls <- generatedCaseClasses(parent, c).find(_.contains(HasAvroAnnotationSuper)).toSeq
+          v <- getApplyPropsSignature(cls)
+        } yield s"def $v = ???"
+        CaseClassFunctions ++ fields
       case _ => Seq.empty
     }
 

diff --git a/src/main/scala/com/spotify/scio/BigQueryTypeInjector.scala b/src/main/scala/com/spotify/scio/BigQueryTypeInjector.scala
@@ -39,8 +39,6 @@ object BigQueryTypeInjector {
     s"$BQTNamespace.toTable"
   )
 
-  private val CaseClassSuper =
-    "_root_.scala.Product"
   private val HasAnnotationSuper =
     "_root_.com.spotify.scio.bigquery.types.BigQueryType.HasAnnotation"
 
@@ -78,14 +76,12 @@ object BigQueryTypeInjector {
     annotation match {
       case a if a.contains(FromQuery) =>
         val simple = """
-          |def query: _root_.java.lang.String = ???
           |def queryRaw: _root_.java.lang.String = ???
           |""".stripMargin
 
         bqQuerySignature(c)
           .map { params =>
             simple + s"""
-              |def query($params): _root_.java.lang.String = ???
               |def queryAsSource($params): _root_.com.spotify.scio.bigquery.Query = ???
               |""".stripMargin
           }
@@ -112,20 +108,16 @@ final class BigQueryTypeInjector extends AnnotationTypeInjector {
   override def injectFunctions(source: ScTypeDefinition): Seq[String] =
     source match {
       case c: ScClass if bqAnnotation(c).isDefined =>
-        val result = for {
-          cc <- Option(c.containingClass)
-          qn <- Option(cc.getQualifiedName)
+        val fields = for {
+          cc <- Option(c.containingClass).toSeq
+          qn <- Option(cc.getQualifiedName).toSeq
           parent = qn.init
-          defs <- {
-            generatedCaseClasses(parent, c)
-              .find(_.contains(HasAnnotationSuper))
-              .map(getApplyPropsSignature)
-              .map(v => s"def $v = ???")
-          }
-        } yield defs
-
-        result.toSeq
-      case _ => Seq.empty
+          cls <- generatedCaseClasses(parent, c).find(_.contains(HasAnnotationSuper)).toSeq
+          v <- getApplyPropsSignature(cls)
+        } yield s"def $v = ???"
+        CaseClassFunctions ++ fields
+      case _ =>
+        Seq.empty
     }
 
   override def injectSupers(source: ScTypeDefinition): Seq[String] =