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 a813414..50589f8 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 @@ -25,6 +25,11 @@ object TracedTransactor { trace(ServiceAndResource(s"$service-db", DefaultResourceName), kleisliTransactor) } + private val commentNamedQueryRegEx = """--\s*Name:\s*(\w+)""".r + + private def extractQueryNameOrSql(sql: String): String = + commentNamedQueryRegEx.findFirstMatchIn(sql).flatMap(m => Option(m.group(1))).getOrElse(sql) + private def formatQuery(q: String): String = q.replace("\n", " ").replaceAll("\\s+", " ").trim() @@ -50,7 +55,7 @@ object TracedTransactor { def runTraced[A](f: TracedOp[A]): TracedOp[A] = Kleisli { case TracedStatement(p, sql) => - Trace[F].span(config.fullyQualifiedSpanName(formatQuery(sql)))( + Trace[F].span(config.fullyQualifiedSpanName(formatQuery(extractQueryNameOrSql(sql))))( Trace[F].put("span.type" -> "db") >> f(p) ) case a => 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 f973e96..8989f32 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,6 +49,15 @@ class TracedTransactorTest extends CatsEffectSuite { ) } + database.test("Trace queries with commented name") { db => + assertIO( + run(sql"""-- Name: selectOne + SELECT 1 WHERE true = ${true: Boolean} + """.query[Int].unique.transact(db)).map(_.last), + SpanData("test-db:db.execute:selectOne", Map("span.type" -> "db")) + ) + } + database.test("Trace updates") { db => case class Test(name: String, age: Int) val create = sql"CREATE TABLE a (id INT, name VARCHAR)".update.run @@ -58,4 +67,17 @@ class TracedTransactorTest extends CatsEffectSuite { SpanData("test-db:db.execute:INSERT INTO a VALUES (?, ?)", Map("span.type" -> "db")) ) } + + database.test("Trace updates with commented name") { db => + case class Test(name: String, age: Int) + val create = sql"CREATE TABLE a (id INT, name VARCHAR)".update.run + val insert = + sql"""-- Name: createNewA + INSERT INTO a VALUES (${2: Int}, ${"abc": String}) + """.update.run + assertIO( + run((create >> insert).transact(db)).map(_.last), + SpanData("test-db:db.execute:createNewA", Map("span.type" -> "db")) + ) + } }