diff --git a/natchez-extras-doobie/src/main/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactor.scala b/natchez-extras-doobie/src/main/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactor.scala index 090de4e..e8bc9ed 100644 --- a/natchez-extras-doobie/src/main/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactor.scala +++ b/natchez-extras-doobie/src/main/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactor.scala @@ -1,7 +1,7 @@ package com.ovoenergy.natchez.extras.doobie import cats.data.Kleisli -import cats.effect.Async +import cats.effect.{Async, Resource} import cats.implicits.catsSyntaxFlatMapOps import com.ovoenergy.natchez.extras.core.Config import com.ovoenergy.natchez.extras.core.Config.ServiceAndResource @@ -54,7 +54,11 @@ object TracedTransactor { ): Transactor[F] = transactor .copy( - interpret0 = createInterpreter(config, Async[F], logHandler).ConnectionInterpreter + interpret0 = createInterpreter(config, Async[F], logHandler).ConnectionInterpreter, + connect0 = in => + Trace[Resource[F, *]].span(config.fullyQualifiedSpanName("connect"))( + Trace[Resource[F, *]].put("span.type" -> "db") >> transactor.connect(in) + ) ) private def createInterpreter[F[_]: Trace]( @@ -101,6 +105,12 @@ object TracedTransactor { super.prepareStatement(a).map(TracedStatement(_, a): PreparedStatement) override def getTypeMap: Nothing = super.getTypeMap.asInstanceOf // See: https://github.com/tpolecat/doobie/blob/v1.0.0-RC4/modules/core/src/test/scala/doobie/util/StrategySuite.scala#L47 + override def commit: Kleisli[F, Connection, Unit] = + Kleisli { c => + Trace[F].span(config.fullyQualifiedSpanName("commit"))( + Trace[F].put("span.type" -> "db") >> super.commit(c) + ) + } } } } diff --git a/natchez-extras-doobie/src/test/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactorTest.scala b/natchez-extras-doobie/src/test/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactorTest.scala index 06af640..d326402 100644 --- a/natchez-extras-doobie/src/test/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactorTest.scala +++ b/natchez-extras-doobie/src/test/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactorTest.scala @@ -49,7 +49,7 @@ class TracedTransactorTest extends CatsEffectSuite { database.test("Trace queries") { db => assertIO( - run(sql"SELECT 1 WHERE true = ${true: Boolean}".query[Int].unique.transact(db)).map(_.last), + run(sql"SELECT 1 WHERE true = ${true: Boolean}".query[Int].unique.transact(db)).map(_.init.last), SpanData("test-db:db.execute:SELECT 1 WHERE true = ?", Map("span.type" -> "db")) ) } @@ -58,7 +58,7 @@ class TracedTransactorTest extends CatsEffectSuite { assertIO( run(sql"""-- Name: selectOne SELECT 1 WHERE true = ${true: Boolean} - """.query[Int].unique.transact(db)).map(_.last), + """.query[Int].unique.transact(db)).map(_.init.last), SpanData("test-db:db.execute:selectOne", Map("span.type" -> "db")) ) } @@ -67,7 +67,7 @@ class TracedTransactorTest extends CatsEffectSuite { val create = sql"CREATE TABLE a (id INT, name VARCHAR)".update.run val insert = sql"INSERT INTO a VALUES (${2: Int}, ${"abc": String})".update.run assertIO( - run((create >> insert).transact(db)).map(_.last), + run((create >> insert).transact(db)).map(_.init.last), SpanData("test-db:db.execute:INSERT INTO a VALUES (?, ?)", Map("span.type" -> "db")) ) } @@ -79,8 +79,28 @@ class TracedTransactorTest extends CatsEffectSuite { INSERT INTO a VALUES (${2: Int}, ${"abc": String}) """.update.run assertIO( - run((create >> insert).transact(db)).map(_.last), + run((create >> insert).transact(db)).map(_.init.last), SpanData("test-db:db.execute:createNewA", Map("span.type" -> "db")) ) } + + database.test("Trace connect") { db => + val create = sql"CREATE TABLE a (id INT, name VARCHAR)".update.run + val insert = + sql"""INSERT INTO a VALUES (${2: Int}, ${"abc": String})""".update.run + assertIO( + run((create >> insert).transact(db)).map(_.drop(1).head), + SpanData("test-db:db.execute:connect", Map("span.type" -> "db")) + ) + } + + database.test("Trace commit") { db => + val create = sql"CREATE TABLE a (id INT, name VARCHAR)".update.run + val insert = + sql"""INSERT INTO a VALUES (${2: Int}, ${"abc": String})""".update.run + assertIO( + run((create >> insert).transact(db)).map(_.last), + SpanData("test-db:db.execute:commit", Map("span.type" -> "db")) + ) + } }