Skip to content

Instantly share code, notes, and snippets.

@tkrs
Last active September 17, 2021 21:51
Show Gist options
  • Save tkrs/b7be4ebca29576c1456b6bfd1d930585 to your computer and use it in GitHub Desktop.
Save tkrs/b7be4ebca29576c1456b6bfd1d930585 to your computer and use it in GitHub Desktop.
Shapeless automatically derives BigQuery TableSchema from case class
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) }
}
@tkrs
Copy link
Author

tkrs commented Feb 16, 2021

GenericData{classInfo=[categories, description, fields, mode, name, policyTags, type], {mode=REQUIRED, name=foo, type=INTEGER}}
GenericData{classInfo=[categories, description, fields, mode, name, policyTags, type], {mode=REQUIRED, name=y, type=STRING}}
GenericData{classInfo=[categories, description, fields, mode, name, policyTags, type], {mode=REQUIRED, name=z, type=TIMESTAMP}}
GenericData{classInfo=[categories, description, fields, mode, name, policyTags, type], {mode=NULLABLE, name=o, type=INTEGER}}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment