Only add tracing to queries over the MySQL protocol

Signed-off-by: Andres Taylor <antaylor@squareup.com>
cashapp · Sep 10, 2019 · 989fab3 · 989fab3
1 parent 9713983
commit 989fab3
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/misk-hibernate-testing/src/test/kotlin/misk/jdbc/JaegerSpanInjectorTest.kt b/misk-hibernate-testing/src/test/kotlin/misk/jdbc/JaegerSpanInjectorTest.kt
@@ -2,8 +2,10 @@ package misk.jdbc
 
 import net.ttddyy.dsproxy.transform.TransformInfo
 import org.junit.jupiter.api.Test
+import org.mockito.Mockito
 import io.opentracing.mock.MockTracer
 import org.assertj.core.api.Assertions.assertThat
+import javax.sql.DataSource
 
 class JaegerSpanInjectorTest {
   @Test
@@ -23,4 +25,26 @@ class JaegerSpanInjectorTest {
     assertThat(result).isEqualTo("/*VT_SPAN_CONTEXT=$contextString*/$query")
     scope.close()
   }
+
+  @Test
+  fun testNotDecorateIfItsVitess() {
+    val tracer = MockTracer()
+    val config = DataSourceConfig(DataSourceType.VITESS)
+    val injector = JaegerSpanInjector(tracer, config)
+    val ds = Mockito.mock(DataSource::class.java)
+
+    assertThat(injector.decorate(ds)).
+        isSameAs(ds)
+  }
+
+  @Test
+  fun testDecoratesIfItsVitessMysql() {
+    val tracer = MockTracer()
+    val config = DataSourceConfig(DataSourceType.VITESS_MYSQL)
+    val injector = JaegerSpanInjector(tracer, config)
+    val ds = Mockito.mock(DataSource::class.java)
+
+    assertThat(injector.decorate(ds)).
+        isNotSameAs(ds)
+  }
 }
diff --git a/misk-hibernate/src/main/kotlin/misk/jdbc/JaegerSpanInjector.kt b/misk-hibernate/src/main/kotlin/misk/jdbc/JaegerSpanInjector.kt
@@ -14,7 +14,7 @@ class JaegerSpanInjector(
   val config: DataSourceConfig
 ) : QueryTransformer, DataSourceDecorator {
     override fun decorate(dataSource: DataSource): DataSource {
-        if (config.type != DataSourceType.VITESS || tracer == null) return dataSource
+        if (config.type != DataSourceType.VITESS_MYSQL || tracer == null) return dataSource
         return ProxyDataSourceBuilder(dataSource).queryTransformer(this).build()
     }