Add trace spans for commit and connect

ovotech · Oct 17, 2023 · 14c3f42 · 14c3f42
1 parent 89efeb5
commit 14c3f42
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 6 deletions.
diff --git a/...z-extras-doobie/src/main/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactor.scala b/...z-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/...tras-doobie/src/test/scala/com/ovoenergy/natchez/extras/doobie/TracedTransactorTest.scala b/...tras-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"))
+    )
+  }
 }