From f45293e8fd3ab239104a6e2c18db4c91f4698ec6 Mon Sep 17 00:00:00 2001
From: Gavin Ray <ray.gavin97@gmail.com>
Date: Wed, 21 Aug 2024 15:36:17 -0400
Subject: [PATCH] Fix OTEL traces not including JDBC statement SQL

---
 .../hasura/ndc/app/application/Application.kt   |  7 +++++++
 .../ndc/app/services/AgroalDataSourceService.kt | 17 ++++-------------
 .../dataConnectors/BaseDataConnectorService.kt  |  2 +-
 3 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/ndc-app/src/main/kotlin/io/hasura/ndc/app/application/Application.kt b/ndc-app/src/main/kotlin/io/hasura/ndc/app/application/Application.kt
index 55a7837..f50b84b 100644
--- a/ndc-app/src/main/kotlin/io/hasura/ndc/app/application/Application.kt
+++ b/ndc-app/src/main/kotlin/io/hasura/ndc/app/application/Application.kt
@@ -15,9 +15,16 @@ object Application {
 
     class MyApp : QuarkusApplication {
         override fun run(vararg args: String): Int {
+            disableOtelJdbcSanitizer()
             Quarkus.waitForExit()
             return 0
         }
+
+        // Disables stripping of whitespace and replacement of literals in SQL queries with "?"
+        // See: https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/6610
+        private fun disableOtelJdbcSanitizer() {
+            System.setProperty("otel.instrumentation.common.db-statement-sanitizer.enabled", "false")
+        }
     }
 
 }
diff --git a/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/AgroalDataSourceService.kt b/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/AgroalDataSourceService.kt
index e030290..24a1f62 100644
--- a/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/AgroalDataSourceService.kt
+++ b/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/AgroalDataSourceService.kt
@@ -9,6 +9,8 @@ import io.agroal.api.security.SimplePassword
 import io.hasura.ndc.common.ConnectorConfiguration
 import io.opentelemetry.instrumentation.annotations.WithSpan
 import io.opentelemetry.instrumentation.jdbc.datasource.OpenTelemetryDataSource
+import io.quarkus.agroal.runtime.AgroalOpenTelemetryWrapper
+import io.quarkus.agroal.runtime.OpenTelemetryAgroalDataSource
 import io.smallrye.config.ConfigMapping
 import io.smallrye.config.WithDefault
 import io.smallrye.config.WithName
@@ -140,7 +142,7 @@ class AgroalDataSourceService {
 
     fun createDataSource(
         connConfig: ConnectorConfiguration
-    ): DataSource {
+    ): AgroalDataSource {
         val configSupplier = mkAgroalDataSourceConfigurationSupplier(connConfig.jdbcUrl, connConfig.jdbcProperties)
         val ds = AgroalDataSource.from(configSupplier).apply {
             loginTimeout = config.connectionFactoryConfiguration().loginTimeout().toSeconds().toInt()
@@ -154,18 +156,7 @@ class AgroalDataSourceService {
         connConfig: ConnectorConfiguration
     ): DataSource {
         val agroalDs = createDataSource(connConfig)
-        return OpenTelemetryDataSource(agroalDs)
-    }
-
-    @WithSpan
-    fun createDataSourceFromConnInfo(
-        connConfig: ConnectorConfiguration,
-        tracing: Boolean = true
-    ): DataSource {
-        return when (tracing) {
-            true -> createTracingDataSource(connConfig)
-            false -> createDataSource(connConfig)
-        }
+        return OpenTelemetryAgroalDataSource(agroalDs)
     }
 
     private object PoolInterceptor : AgroalPoolInterceptor {
diff --git a/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/dataConnectors/BaseDataConnectorService.kt b/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/dataConnectors/BaseDataConnectorService.kt
index 6d5d341..b072360 100644
--- a/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/dataConnectors/BaseDataConnectorService.kt
+++ b/ndc-app/src/main/kotlin/io/hasura/ndc/app/services/dataConnectors/BaseDataConnectorService.kt
@@ -38,7 +38,7 @@ class JDBCDataSourceProvider : IDataSourceProvider {
 
     override fun getDataSource(config: ConnectorConfiguration): DataSource {
         if (dataSource == null) {
-            dataSource = agroalDataSourceService.createDataSource(config)
+            dataSource = agroalDataSourceService.createTracingDataSource(config)
         }
         return dataSource!!
     }