Support ObjectId in schema-bson

.github/workflows/ci.yml

@@ -33,7 +33,7 @@ jobs:
 
  build:
  runs-on: ubuntu-20.04
- timeout-minutes: 30
+ timeout-minutes: 40
  strategy:
  fail-fast: false
  matrix:
@@ -71,4 +71,4 @@ jobs:
  PGP_SECRET: ${{ secrets.PGP_SECRET }}
  SONATYPE_PASSWORD: ${{ secrets.SONATYPE_PASSWORD }}
  SONATYPE_USERNAME: ${{ secrets.SONATYPE_USERNAME }}
-
+

zio-schema-bson/src/main/scala/zio/schema/codec/BsonSchemaCodec.scala

@@ -7,6 +7,7 @@ import scala.collection.compat._
 import scala.collection.immutable.{ HashMap, ListMap }
 import scala.jdk.CollectionConverters._
 
+import org.bson.types.ObjectId
 import org.bson.{ BsonDocument, BsonNull, BsonReader, BsonType, BsonValue, BsonWriter }
 
 import zio.bson.BsonBuilder._
@@ -505,6 +506,10 @@ object BsonSchemaCodec {
  directEncoder =>
  override def encode(writer: BsonWriter, value: DynamicValue, ctx: BsonEncoder.EncoderContext): Unit =
  value match {
+ case DynamicValue.Record(_, values) if values.headOption.exists(_._1 == bson.ObjectIdTag) =>
+ val id = values.head._2.toTypedValueOption[String].get
+ writer.writeObjectId(new ObjectId(id))
+
  case DynamicValue.Record(_, values) =>
  val nextCtx = BsonEncoder.EncoderContext.default
 
@@ -547,6 +552,11 @@ object BsonSchemaCodec {
 
  override def toBsonValue(value: DynamicValue): BsonValue =
  value match {
+ case DynamicValue.Record(_, values) if values.headOption.exists(_._1 == bson.ObjectIdTag) =>
+ val id = values.head._2.toTypedValueOption[String].get
+ val objectId = new ObjectId(id)
+ objectId.toBsonValue
+
  case DynamicValue.Record(_, values) =>
  new BsonDocument(values.view.map {
  case (key, value) => element(key, directEncoder.toBsonValue(value))
@@ -835,7 +845,15 @@ object BsonSchemaCodec {
  DynamicValue.Primitive(Chunk.fromArray(bsonValue.asBinary().getData), StandardType.BinaryType)
  case BsonType.UNDEFINED => DynamicValue.NoneValue
  case BsonType.OBJECT_ID =>
- DynamicValue.Primitive(bsonValue.asObjectId().getValue.toHexString, StandardType.StringType)
+ DynamicValue.Record(
+ TypeId.Structural,
+ ListMap(
+ bson.ObjectIdTag -> DynamicValue.Primitive(
+ bsonValue.asObjectId().getValue.toHexString,
+ StandardType.StringType
+ )
+ )
+ )
  case BsonType.BOOLEAN => DynamicValue.Primitive(bsonValue.asBoolean().getValue, StandardType.BoolType)
  case BsonType.DATE_TIME =>
  DynamicValue.Primitive(Instant.ofEpochMilli(bsonValue.asDateTime().getValue), StandardType.InstantType)
@@ -1146,39 +1164,49 @@ object BsonSchemaCodec {
 
  private val len = nonTransientFields.length
 
- def encode(writer: BsonWriter, value: Z, ctx: BsonEncoder.EncoderContext): Unit = {
- val nextCtx = ctx.copy(inlineNextObject = false)
+ def encode(writer: BsonWriter, value: Z, ctx: BsonEncoder.EncoderContext): Unit =
+ if (names.size == 1 && names(0) == bson.ObjectIdTag) {
+ val fieldValue = nonTransientFields(0).get(value)
+ val id = new ObjectId(fieldValue.toString)
+ writer.writeObjectId(id)
+ } else {
+ val nextCtx = ctx.copy(inlineNextObject = false)
+
+ if (!ctx.inlineNextObject) writer.writeStartDocument()
 
- if (!ctx.inlineNextObject) writer.writeStartDocument()
+  var i = 0
 
- var i = 0
+ while (i < len) {
+ val tc = tcs(i)
+ val fieldValue = nonTransientFields(i).get(value)
 
- while (i < len) {
- val tc = tcs(i)
- val fieldValue = nonTransientFields(i).get(value)
+ if (keepNulls || !tc.isAbsent(fieldValue)) {
+ writer.writeName(names(i))
+ tc.encode(writer, fieldValue, nextCtx)
+ }
 
- if (keepNulls || !tc.isAbsent(fieldValue)) {
- writer.writeName(names(i))
- tc.encode(writer, fieldValue, nextCtx)
+ i += 1
  }
 
- i += 1
+ if (!ctx.inlineNextObject) writer.writeEndDocument()
  }
 
- if (!ctx.inlineNextObject) writer.writeEndDocument()
- }
-
- def toBsonValue(value: Z): BsonValue = {
- val elements = nonTransientFields.indices.view.flatMap { idx =>
- val fieldValue = nonTransientFields(idx).get(value)
- val tc = tcs(idx)
+ def toBsonValue(value: Z): BsonValue =
+ if (names.size == 1 && names(0) == bson.ObjectIdTag) {
+ val fieldValue = nonTransientFields(0).get(value)
+ val id = new ObjectId(fieldValue.toString)
+ id.toBsonValue
+ } else {
+ val elements = nonTransientFields.indices.view.flatMap { idx =>
+ val fieldValue = nonTransientFields(idx).get(value)
+ val tc = tcs(idx)
 
- if (keepNulls || !tc.isAbsent(fieldValue)) Some(element(names(idx), tc.toBsonValue(fieldValue)))
- else None
- }.to(Chunk)
+  if (keepNulls || !tc.isAbsent(fieldValue)) Some(element(names(idx), tc.toBsonValue(fieldValue)))
+  else None
+  }.to(Chunk)
 
- new BsonDocument(elements.asJava)
- }
+  new BsonDocument(elements.asJava)
+  }
  }
  }
 
@@ -1190,7 +1218,6 @@ object BsonSchemaCodec {
  private[codec] def caseClassDecoder[Z](caseClassSchema: Schema.Record[Z]): BsonDecoder[Z] = {
  val fields = caseClassSchema.fields
  val len: Int = fields.length
- Array.ofDim[Any](len)
  val fieldNames = fields.map { f =>
  f.annotations.collectFirst { case bsonField(n) => n }.getOrElse(f.name.asInstanceOf[String])
  }.toArray
@@ -1210,91 +1237,112 @@ object BsonSchemaCodec {
  lazy val tcs: Array[BsonDecoder[Any]] = schemas.map(s => schemaDecoder(s).asInstanceOf[BsonDecoder[Any]])
 
  new BsonDecoder[Z] {
- def decodeUnsafe(reader: BsonReader, trace: List[BsonTrace], ctx: BsonDecoder.BsonDecoderContext): Z = unsafeCall(trace) {
- reader.readStartDocument()
-
- val nextCtx = BsonDecoder.BsonDecoderContext.default
- val ps: Array[Any] = Array.ofDim(len)
-
- while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
- val name = reader.readName()
- val idx = indexes.getOrElse(name, -1)
-
- if (idx >= 0) {
- val nextTrace = spans(idx) :: trace
- val tc = tcs(idx)
- if (ps(idx) != null) throw BsonDecoder.Error(nextTrace, "duplicate")
- ps(idx) = if ((fields(idx).optional || fields(idx).transient) && fields(idx).defaultValue.isDefined) {
- val opt = BsonDecoder.option(tc).decodeUnsafe(reader, nextTrace, nextCtx)
- opt.getOrElse(fields(idx).defaultValue.get)
- } else {
- tc.decodeUnsafe(reader, nextTrace, nextCtx)
+ def decodeUnsafe(reader: BsonReader, trace: List[BsonTrace], ctx: BsonDecoder.BsonDecoderContext): Z =
+ if (fieldNames.size == 1 && fieldNames(0) == bson.ObjectIdTag) {
+ val id = reader.readObjectId.toHexString
+ Unsafe.unsafe { implicit u =>
+ caseClassSchema.construct(Chunk.fromArray(Array(id))) match {
+ case Left(err) => throw BsonDecoder.Error(trace, s"Failed to construct case class: $err")
+ case Right(value) => value
+ }
+ }
+ } else {
+ unsafeCall(trace) {
+ reader.readStartDocument()
+
+ val nextCtx = BsonDecoder.BsonDecoderContext.default
+ val ps: Array[Any] = Array.ofDim(len)
+
+ while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
+ val name = reader.readName()
+ val idx = indexes.getOrElse(name, -1)
+
+ if (idx >= 0) {
+ val nextTrace = spans(idx) :: trace
+ val tc = tcs(idx)
+ if (ps(idx) != null) throw BsonDecoder.Error(nextTrace, "duplicate")
+ ps(idx) = if ((fields(idx).optional || fields(idx).transient) && fields(idx).defaultValue.isDefined) {
+ val opt = BsonDecoder.option(tc).decodeUnsafe(reader, nextTrace, nextCtx)
+ opt.getOrElse(fields(idx).defaultValue.get)
+ } else {
+ tc.decodeUnsafe(reader, nextTrace, nextCtx)
+ }
+ } else if (noExtra && !ctx.ignoreExtraField.contains(name)) {
+ throw BsonDecoder.Error(BsonTrace.Field(name) :: trace, "Invalid extra field.")
+ } else reader.skipValue()
  }
- } else if (noExtra && !ctx.ignoreExtraField.contains(name)) {
- throw BsonDecoder.Error(BsonTrace.Field(name) :: trace, "Invalid extra field.")
- } else reader.skipValue()
- }
 
- var i = 0
- while (i < len) {
- if (ps(i) == null) {
- if ((fields(i).optional || fields(i).transient) && fields(i).defaultValue.isDefined) {
- ps(i) = fields(i).defaultValue.get
- } else {
- ps(i) = tcs(i).decodeMissingUnsafe(spans(i) :: trace)
+ var i = 0
+ while (i < len) {
+ if (ps(i) == null) {
+ if ((fields(i).optional || fields(i).transient) && fields(i).defaultValue.isDefined) {
+ ps(i) = fields(i).defaultValue.get
+ } else {
+ ps(i) = tcs(i).decodeMissingUnsafe(spans(i) :: trace)
+ }
+ }
+ i += 1
  }
- }
- i += 1
- }
 
- reader.readEndDocument()
+  reader.readEndDocument()
 
- Unsafe.unsafe { implicit u =>
- caseClassSchema.construct(Chunk.fromArray(ps)) match {
- case Left(err) => throw BsonDecoder.Error(trace, s"Failed to construct case class: $err")
- case Right(value) => value
+ Unsafe.unsafe { implicit u =>
+ caseClassSchema.construct(Chunk.fromArray(ps)) match {
+ case Left(err) => throw BsonDecoder.Error(trace, s"Failed to construct case class: $err")
+ case Right(value) => value
+ }
+ }
  }
  }
- }
 
  def fromBsonValueUnsafe(value: BsonValue, trace: List[BsonTrace], ctx: BsonDecoder.BsonDecoderContext): Z =
- assumeType(trace)(BsonType.DOCUMENT, value) { value =>
- val nextCtx = BsonDecoder.BsonDecoderContext.default
- val ps: Array[Any] = Array.ofDim(len)
-
- value.asDocument().asScala.foreachEntry { (name, value) =>
- val idx = indexes.getOrElse(name, -1)
-
- if (idx >= 0) {
- val nextTrace = spans(idx) :: trace
- val tc = tcs(idx)
- if (ps(idx) != null) throw BsonDecoder.Error(nextTrace, "duplicate")
- ps(idx) = if ((fields(idx).optional || fields(idx).transient) && fields(idx).defaultValue.isDefined) {
- val opt = BsonDecoder.option(tc).fromBsonValueUnsafe(value, nextTrace, nextCtx)
- opt.getOrElse(fields(idx).defaultValue.get)
- } else {
- tc.fromBsonValueUnsafe(value, nextTrace, nextCtx)
- }
- } else if (noExtra && !ctx.ignoreExtraField.contains(name))
- throw BsonDecoder.Error(BsonTrace.Field(name) :: trace, "Invalid extra field.")
+ if (value.getBsonType == BsonType.OBJECT_ID) {
+ Unsafe.unsafe { implicit u =>
+ val ps: Array[Any] = Array(value.asObjectId.getValue.toHexString)
+ caseClassSchema.construct(Chunk.fromArray(ps)) match {
+ case Left(err) => throw BsonDecoder.Error(trace, s"Failed to construct case class: $err")
+ case Right(value) => value
+ }
  }
+ } else {
+ assumeType(trace)(BsonType.DOCUMENT, value) { value =>
+ val nextCtx = BsonDecoder.BsonDecoderContext.default
+ val ps: Array[Any] = Array.ofDim(len)
+
+ value.asDocument().asScala.foreachEntry { (name, value) =>
+ val idx = indexes.getOrElse(name, -1)
+
+ if (idx >= 0) {
+ val nextTrace = spans(idx) :: trace
+ val tc = tcs(idx)
+ if (ps(idx) != null) throw BsonDecoder.Error(nextTrace, "duplicate")
+ ps(idx) = if ((fields(idx).optional || fields(idx).transient) && fields(idx).defaultValue.isDefined) {
+ val opt = BsonDecoder.option(tc).fromBsonValueUnsafe(value, nextTrace, nextCtx)
+ opt.getOrElse(fields(idx).defaultValue.get)
+ } else {
+ tc.fromBsonValueUnsafe(value, nextTrace, nextCtx)
+ }
+ } else if (noExtra && !ctx.ignoreExtraField.contains(name))
+ throw BsonDecoder.Error(BsonTrace.Field(name) :: trace, "Invalid extra field.")
+ }
 
- var i = 0
- while (i < len) {
- if (ps(i) == null) {
- ps(i) = if ((fields(i).optional || fields(i).transient) && fields(i).defaultValue.isDefined) {
- fields(i).defaultValue.get
- } else {
- tcs(i).decodeMissingUnsafe(spans(i) :: trace)
+ var i = 0
+ while (i < len) {
+ if (ps(i) == null) {
+ ps(i) = if ((fields(i).optional || fields(i).transient) && fields(i).defaultValue.isDefined) {
+ fields(i).defaultValue.get
+ } else {
+ tcs(i).decodeMissingUnsafe(spans(i) :: trace)
+ }
  }
+ i += 1
  }
- i += 1
- }
 
- Unsafe.unsafe { implicit u =>
- caseClassSchema.construct(Chunk.fromArray(ps)) match {
- case Left(err) => throw BsonDecoder.Error(trace, s"Failed to construct case class: $err")
- case Right(value) => value
+ Unsafe.unsafe { implicit u =>
+ caseClassSchema.construct(Chunk.fromArray(ps)) match {
+ case Left(err) => throw BsonDecoder.Error(trace, s"Failed to construct case class: $err")
+ case Right(value) => value
+ }
  }
  }
  }

zio-schema-bson/src/main/scala/zio/schema/codec/package.scala

@@ -0,0 +1,22 @@
+package zio.schema.codec
+
+import org.bson.types.ObjectId
+
+import zio.schema.{ Schema, TypeId }
+
+package object bson {
+ val ObjectIdTag = "$oid"
+
+ implicit val ObjectIdSchema: Schema[ObjectId] =
+ Schema.CaseClass1[String, ObjectId](
+ id0 = TypeId.fromTypeName("ObjectId"),
+ field0 = Schema.Field(
+ name0 = ObjectIdTag,
+ schema0 = Schema[String],
+ get0 = _.toHexString,
+ set0 = (_, idStr) => new ObjectId(idStr)
+ ),
+ defaultConstruct0 = new ObjectId(_)
+ )
+
+}