Last active
September 17, 2021 21:51
-
-
Save tkrs/b7be4ebca29576c1456b6bfd1d930585 to your computer and use it in GitHub Desktop.
Shapeless automatically derives BigQuery TableSchema from case class
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import com.google.api.services.bigquery.model.TableFieldSchema | |
import com.google.api.services.bigquery.model.TableSchema | |
import org.joda.time.Instant | |
import shapeless._ | |
import shapeless.ops.hlist.FillWith | |
import shapeless.ops.hlist.Mapper | |
import shapeless.ops.hlist.ToList | |
import shapeless.ops.record.Keys | |
import shapeless.ops.record.Values | |
import scala.annotation.StaticAnnotation | |
import scala.annotation.nowarn | |
import scala.jdk.CollectionConverters._ | |
object Main extends App { | |
sealed abstract class Type(val value: String) | |
object Type { | |
case object Timestamp extends Type("TIMESTAMP") | |
case object String extends Type("STRING") | |
case object Integer extends Type("INTEGER") | |
} | |
sealed abstract class Mode(val value: String) | |
object Mode { | |
case object Required extends Mode("REQUIRED") | |
case object Nullable extends Mode("NULLABLE") | |
case object Repeated extends Mode("REPEATED") | |
} | |
case class Schema() extends StaticAnnotation | |
case class TableField( | |
name: String = null, | |
`type`: Option[Type] = None, | |
mode: Option[Mode] = None, | |
categories: List[String] = Nil, | |
description: Option[String] = None, | |
policyTags: List[String] = Nil | |
) extends StaticAnnotation | |
trait ToSchema[A] { | |
def apply(): TableSchema | |
} | |
object ToSchema { | |
implicit def apply[A](implicit A: ToSchema[A]): ToSchema[A] = A | |
object toNull extends Poly0 { | |
implicit def default[T]: ProductCase.Aux[HNil, T] = at[T](null.asInstanceOf[T]) | |
} | |
trait ToType[A] { | |
def apply(): Type | |
} | |
object ToType { | |
implicit def apply[A](implicit A: ToType[A]): ToType[A] = A | |
implicit val string: ToType[String] = () => Type.String | |
implicit val long: ToType[Long] = () => Type.Integer | |
implicit val instant: ToType[Instant] = () => Type.Timestamp | |
} | |
object toTypeMode extends Poly1 { | |
implicit def default[A](implicit A: ToType[A]): Case.Aux[A, (Type, Mode)] = at(_ => (A(), Mode.Required)) | |
implicit def option[A](implicit A: ToType[A]): Case.Aux[Option[A], (Type, Mode)] = at(_ => (A(), Mode.Nullable)) | |
} | |
implicit def cc[A, R <: HList, K <: HList, V <: HList, F <: HList, Out <: HList](implicit | |
@nowarn gen: LabelledGeneric.Aux[A, R], | |
@nowarn schema: Annotation[Schema, A], | |
@nowarn values: Values.Aux[R, V], | |
keys: Keys.Aux[R, K], | |
keysToList: ToList[K, Symbol], | |
tableFields: Annotations.Aux[TableField, A, F], | |
tableFieldsToList: ToList[F, Option[TableField]], | |
fillNull: FillWith[toNull.type, V], | |
mapType: Mapper.Aux[toTypeMode.type, V, Out], | |
toTypeList: ToList[Out, (Type, Mode)] | |
): ToSchema[A] = { | |
val schema = keysToList(keys()) | |
.map(_.name) | |
.zip(toTypeList(mapType(fillNull()))) | |
.zip(tableFieldsToList(tableFields())) | |
() => | |
new TableSchema().setFields( | |
schema.map { case ((name, (t, m)), f) => | |
val _s = new TableFieldSchema() | |
.setName(f.flatMap(a => Option(a.name)).getOrElse(name)) | |
.setType(f.flatMap(_.`type`).getOrElse(t).value) | |
.setMode(f.flatMap(_.mode).getOrElse(m).value) | |
f.flatMap(_.description).foreach(_s.setDescription) | |
f.flatMap(a => Option.when(a.categories.nonEmpty)(a.categories.asJava)) | |
.map(new TableFieldSchema.Categories().setNames) | |
.foreach(_s.setCategories) | |
f.flatMap(a => Option.when(a.categories.nonEmpty)(a.policyTags.asJava)) | |
.map(new TableFieldSchema.PolicyTags().setNames) | |
.foreach(_s.setPolicyTags) | |
_s | |
}.asJava | |
) | |
} | |
} | |
@Schema() | |
case class Foo( | |
@TableField("foo") x: Long, | |
y: String, | |
z: Instant, | |
o: Option[Long] | |
) | |
val toSchema = ToSchema[Foo] | |
toSchema().getFields().forEach { s => println(s) } | |
} |
Author
tkrs
commented
Feb 16, 2021
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment