Merge pull request #114 from ovotech/ajkaanbal/fix-updates-with-gener…

…ated-keys fix: add traces to updates with generared keys
ovotech · Feb 9, 2024 · 9bca66c · 9bca66c
2 parents 9f71f73 + f8232a0
commit 9bca66c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
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
@@ -99,6 +99,20 @@ object TracedTransactor {
         new ConnectionInterpreter {
           override def prepareStatement(a: String): Kleisli[F, Connection, PreparedStatement] =
             super.prepareStatement(a).map(TracedStatement(_, a): PreparedStatement)
+          override def prepareStatement(
+            a: String,
+            b: Array[String]
+          ): Kleisli[F, Connection, PreparedStatement] =
+            super.prepareStatement(a, b).map(TracedStatement(_, a): PreparedStatement)
+          override def prepareStatement(a: String, b: Array[Int]): Kleisli[F, Connection, PreparedStatement] =
+            super.prepareStatement(a, b).map(TracedStatement(_, a): PreparedStatement)
+          override def prepareStatement(a: String, b: Int) =
+            super.prepareStatement(a, b).map(TracedStatement(_, a): PreparedStatement)
+          override def prepareStatement(a: String, b: Int, c: Int) =
+            super.prepareStatement(a, b, c).map(TracedStatement(_, a): PreparedStatement)
+          override def prepareStatement(a: String, b: Int, c: Int, d: Int) =
+            super.prepareStatement(a, b, c, d).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
         }

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
@@ -65,7 +65,18 @@ class TracedTransactorTest extends CatsEffectSuite {
 
   database.test("Trace updates") { 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
+    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:INSERT INTO a VALUES (?, ?)", Map("span.type" -> "db"))
+    )
+  }
+
+  database.test("Trace updates withUniqueGeneratedKeys") { 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.withUniqueGeneratedKeys[String]("name")
     assertIO(
       run((create >> insert).transact(db)).map(_.last),
       SpanData("test-db:db.execute:INSERT INTO a VALUES (?, ?)", Map("span.type" -> "db"))