feat(telemetry): add dynamic "db.system" attribute to OTEL spans (#5019)

* feat(telemetry): add dynamic "db.system" attribute to OTEL spans * chore: clippy * chore(driver-adapters): add "db_system_name()" to "AdapterFlavour" * chore(driver-adapters): rename SYSTEM_NAME to DB_SYSTEM_NAME, and change member variable to constant * chore(driver-adapters): get rid of db_system_name(), adapt name() on Connector trait * chore(driver-adapters): fix typos * chore: clippy * chore: fix compilation on wasm32 * feat(tracing): add "otel.kind" = "client" on query spans * chore: fix Wasm compilation * feat(core): extract otel_kind from span attributes * chore(core): clippy * chore(core): remove percentage sigil for constant strings * chore(core): rename "system_name" to "db_system_name" * feat(core): add mongodb-specific "db.collection.name" and "db.operation.name" attributes in spans * chore(core): review comments * chore: Update query-engine/connectors/mongodb-query-connector/src/root_queries/mod.rs Co-authored-by: Alexey Orlenko <[email protected]> * chore: address review comments --------- Co-authored-by: Alexey Orlenko <[email protected]>
prisma · Oct 22, 2024 · 8263a1f · 8263a1f
1 parent 4fe298b
commit 8263a1f
Show file tree

Hide file tree

Showing 18 changed files with 236 additions and 115 deletions.
diff --git a/quaint/src/connector/metrics.rs b/quaint/src/connector/metrics.rs
@@ -4,12 +4,19 @@ use crate::ast::{Params, Value};
 use crosstarget_utils::time::ElapsedTimeCounter;
 use std::future::Future;
 
-pub async fn query<'a, F, T, U>(tag: &'static str, query: &'a str, params: &'a [Value<'_>], f: F) -> crate::Result<T>
+pub async fn query<'a, F, T, U>(
+    tag: &'static str,
+    db_system_name: &'static str,
+    query: &'a str,
+    params: &'a [Value<'_>],
+    f: F,
+) -> crate::Result<T>
 where
     F: FnOnce() -> U + 'a,
     U: Future<Output = crate::Result<T>>,
 {
-    let span = info_span!("quaint:query", "db.statement" = %query);
+    let span =
+        info_span!("quaint:query", "db.system" = db_system_name, "db.statement" = %query, "otel.kind" = "client");
     do_query(tag, query, params, f).instrument(span).await
 }
 

diff --git a/quaint/src/connector/mssql/native/mod.rs b/quaint/src/connector/mssql/native/mod.rs
@@ -30,6 +30,7 @@ use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt};
 pub use tiberius;
 
 static SQL_SERVER_DEFAULT_ISOLATION: IsolationLevel = IsolationLevel::ReadCommitted;
+const DB_SYSTEM_NAME: &str = "mssql";
 
 #[async_trait]
 impl TransactionCapable for Mssql {
@@ -130,7 +131,7 @@ impl Queryable for Mssql {
     }
 
     async fn query_raw(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<ResultSet> {
-        metrics::query("mssql.query_raw", sql, params, move || async move {
+        metrics::query("mssql.query_raw", DB_SYSTEM_NAME, sql, params, move || async move {
             let mut client = self.client.lock().await;
 
             let mut query = tiberius::Query::new(sql);
@@ -193,7 +194,7 @@ impl Queryable for Mssql {
     }
 
     async fn execute_raw(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<u64> {
-        metrics::query("mssql.execute_raw", sql, params, move || async move {
+        metrics::query("mssql.execute_raw", DB_SYSTEM_NAME, sql, params, move || async move {
             let mut query = tiberius::Query::new(sql);
 
             for param in params {
@@ -213,7 +214,7 @@ impl Queryable for Mssql {
     }
 
     async fn raw_cmd(&self, cmd: &str) -> crate::Result<()> {
-        metrics::query("mssql.raw_cmd", cmd, &[], move || async move {
+        metrics::query("mssql.raw_cmd", DB_SYSTEM_NAME, cmd, &[], move || async move {
             let mut client = self.client.lock().await;
             self.perform_io(client.simple_query(cmd)).await?.into_results().await?;
             Ok(())

diff --git a/quaint/src/connector/mysql/native/mod.rs b/quaint/src/connector/mysql/native/mod.rs
@@ -68,6 +68,8 @@ impl MysqlUrl {
     }
 }
 
+const DB_SYSTEM_NAME: &str = "mysql";
+
 /// A connector interface for the MySQL database.
 #[derive(Debug)]
 pub struct Mysql {
@@ -195,7 +197,7 @@ impl Queryable for Mysql {
     }
 
     async fn query_raw(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<ResultSet> {
-        metrics::query("mysql.query_raw", sql, params, move || async move {
+        metrics::query("mysql.query_raw", DB_SYSTEM_NAME, sql, params, move || async move {
             self.prepared(sql, |stmt| async move {
                 let mut conn = self.conn.lock().await;
                 let rows: Vec<my::Row> = conn.exec(&stmt, conversion::conv_params(params)?).await?;
@@ -280,7 +282,7 @@ impl Queryable for Mysql {
     }
 
     async fn execute_raw(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<u64> {
-        metrics::query("mysql.execute_raw", sql, params, move || async move {
+        metrics::query("mysql.execute_raw", DB_SYSTEM_NAME, sql, params, move || async move {
             self.prepared(sql, |stmt| async move {
                 let mut conn = self.conn.lock().await;
                 conn.exec_drop(stmt, conversion::conv_params(params)?).await?;
@@ -297,7 +299,7 @@ impl Queryable for Mysql {
     }
 
     async fn raw_cmd(&self, cmd: &str) -> crate::Result<()> {
-        metrics::query("mysql.raw_cmd", cmd, &[], move || async move {
+        metrics::query("mysql.raw_cmd", DB_SYSTEM_NAME, cmd, &[], move || async move {
             self.perform_io(|| async move {
                 let mut conn = self.conn.lock().await;
                 let mut result = cmd.run(&mut *conn).await?;

diff --git a/quaint/src/connector/postgres/native/mod.rs b/quaint/src/connector/postgres/native/mod.rs
@@ -55,6 +55,9 @@ impl Debug for PostgresClient {
     }
 }
 
+const DB_SYSTEM_NAME_POSTGRESQL: &str = "postgresql";
+const DB_SYSTEM_NAME_COCKROACHDB: &str = "cockroachdb";
+
 /// A connector interface for the PostgreSQL database.
 #[derive(Debug)]
 pub struct PostgreSql {
@@ -65,6 +68,7 @@ pub struct PostgreSql {
     is_healthy: AtomicBool,
     is_cockroachdb: bool,
     is_materialize: bool,
+    db_system_name: &'static str,
 }
 
 /// Key uniquely representing an SQL statement in the prepared statements cache.
@@ -285,6 +289,12 @@ impl PostgreSql {
             }
         }
 
+        let db_system_name = if is_cockroachdb {
+            DB_SYSTEM_NAME_COCKROACHDB
+        } else {
+            DB_SYSTEM_NAME_POSTGRESQL
+        };
+
         Ok(Self {
             client: PostgresClient(client),
             socket_timeout: url.query_params.socket_timeout,
@@ -293,6 +303,7 @@ impl PostgreSql {
             is_healthy: AtomicBool::new(true),
             is_cockroachdb,
             is_materialize,
+            db_system_name,
         })
     }
 
@@ -308,6 +319,7 @@ impl PostgreSql {
             is_healthy: AtomicBool::new(true),
             is_cockroachdb: false,
             is_materialize: false,
+            db_system_name: DB_SYSTEM_NAME_POSTGRESQL,
         })
     }
 
@@ -539,72 +551,84 @@ impl Queryable for PostgreSql {
     async fn query_raw(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<ResultSet> {
         self.check_bind_variables_len(params)?;
 
-        metrics::query("postgres.query_raw", sql, params, move || async move {
-            let stmt = self.fetch_cached(sql, &[]).await?;
-
-            if stmt.params().len() != params.len() {
-                let kind = ErrorKind::IncorrectNumberOfParameters {
-                    expected: stmt.params().len(),
-                    actual: params.len(),
-                };
-
-                return Err(Error::builder(kind).build());
-            }
+        metrics::query(
+            "postgres.query_raw",
+            self.db_system_name,
+            sql,
+            params,
+            move || async move {
+                let stmt = self.fetch_cached(sql, &[]).await?;
+
+                if stmt.params().len() != params.len() {
+                    let kind = ErrorKind::IncorrectNumberOfParameters {
+                        expected: stmt.params().len(),
+                        actual: params.len(),
+                    };
+
+                    return Err(Error::builder(kind).build());
+                }
 
-            let rows = self
-                .perform_io(self.client.0.query(&stmt, conversion::conv_params(params).as_slice()))
-                .await?;
+                let rows = self
+                    .perform_io(self.client.0.query(&stmt, conversion::conv_params(params).as_slice()))
+                    .await?;
 
-            let col_types = stmt
-                .columns()
-                .iter()
-                .map(|c| PGColumnType::from_pg_type(c.type_()))
-                .map(ColumnType::from)
-                .collect::<Vec<_>>();
-            let mut result = ResultSet::new(stmt.to_column_names(), col_types, Vec::new());
+                let col_types = stmt
+                    .columns()
+                    .iter()
+                    .map(|c| PGColumnType::from_pg_type(c.type_()))
+                    .map(ColumnType::from)
+                    .collect::<Vec<_>>();
+                let mut result = ResultSet::new(stmt.to_column_names(), col_types, Vec::new());
 
-            for row in rows {
-                result.rows.push(row.get_result_row()?);
-            }
+                for row in rows {
+                    result.rows.push(row.get_result_row()?);
+                }
 
-            Ok(result)
-        })
+                Ok(result)
+            },
+        )
         .await
     }
 
     async fn query_raw_typed(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<ResultSet> {
         self.check_bind_variables_len(params)?;
 
-        metrics::query("postgres.query_raw", sql, params, move || async move {
-            let stmt = self.fetch_cached(sql, params).await?;
-
-            if stmt.params().len() != params.len() {
-                let kind = ErrorKind::IncorrectNumberOfParameters {
-                    expected: stmt.params().len(),
-                    actual: params.len(),
-                };
-
-                return Err(Error::builder(kind).build());
-            }
+        metrics::query(
+            "postgres.query_raw",
+            self.db_system_name,
+            sql,
+            params,
+            move || async move {
+                let stmt = self.fetch_cached(sql, params).await?;
+
+                if stmt.params().len() != params.len() {
+                    let kind = ErrorKind::IncorrectNumberOfParameters {
+                        expected: stmt.params().len(),
+                        actual: params.len(),
+                    };
+
+                    return Err(Error::builder(kind).build());
+                }
 
-            let col_types = stmt
-                .columns()
-                .iter()
-                .map(|c| PGColumnType::from_pg_type(c.type_()))
-                .map(ColumnType::from)
-                .collect::<Vec<_>>();
-            let rows = self
-                .perform_io(self.client.0.query(&stmt, conversion::conv_params(params).as_slice()))
-                .await?;
+                let col_types = stmt
+                    .columns()
+                    .iter()
+                    .map(|c| PGColumnType::from_pg_type(c.type_()))
+                    .map(ColumnType::from)
+                    .collect::<Vec<_>>();
+                let rows = self
+                    .perform_io(self.client.0.query(&stmt, conversion::conv_params(params).as_slice()))
+                    .await?;
 
-            let mut result = ResultSet::new(stmt.to_column_names(), col_types, Vec::new());
+                let mut result = ResultSet::new(stmt.to_column_names(), col_types, Vec::new());
 
-            for row in rows {
-                result.rows.push(row.get_result_row()?);
-            }
+                for row in rows {
+                    result.rows.push(row.get_result_row()?);
+                }
 
-            Ok(result)
-        })
+                Ok(result)
+            },
+        )
         .await
     }
 
@@ -692,53 +716,65 @@ impl Queryable for PostgreSql {
     async fn execute_raw(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<u64> {
         self.check_bind_variables_len(params)?;
 
-        metrics::query("postgres.execute_raw", sql, params, move || async move {
-            let stmt = self.fetch_cached(sql, &[]).await?;
-
-            if stmt.params().len() != params.len() {
-                let kind = ErrorKind::IncorrectNumberOfParameters {
-                    expected: stmt.params().len(),
-                    actual: params.len(),
-                };
-
-                return Err(Error::builder(kind).build());
-            }
+        metrics::query(
+            "postgres.execute_raw",
+            self.db_system_name,
+            sql,
+            params,
+            move || async move {
+                let stmt = self.fetch_cached(sql, &[]).await?;
+
+                if stmt.params().len() != params.len() {
+                    let kind = ErrorKind::IncorrectNumberOfParameters {
+                        expected: stmt.params().len(),
+                        actual: params.len(),
+                    };
+
+                    return Err(Error::builder(kind).build());
+                }
 
-            let changes = self
-                .perform_io(self.client.0.execute(&stmt, conversion::conv_params(params).as_slice()))
-                .await?;
+                let changes = self
+                    .perform_io(self.client.0.execute(&stmt, conversion::conv_params(params).as_slice()))
+                    .await?;
 
-            Ok(changes)
-        })
+                Ok(changes)
+            },
+        )
         .await
     }
 
     async fn execute_raw_typed(&self, sql: &str, params: &[Value<'_>]) -> crate::Result<u64> {
         self.check_bind_variables_len(params)?;
 
-        metrics::query("postgres.execute_raw", sql, params, move || async move {
-            let stmt = self.fetch_cached(sql, params).await?;
-
-            if stmt.params().len() != params.len() {
-                let kind = ErrorKind::IncorrectNumberOfParameters {
-                    expected: stmt.params().len(),
-                    actual: params.len(),
-                };
-
-                return Err(Error::builder(kind).build());
-            }
+        metrics::query(
+            "postgres.execute_raw",
+            self.db_system_name,
+            sql,
+            params,
+            move || async move {
+                let stmt = self.fetch_cached(sql, params).await?;
+
+                if stmt.params().len() != params.len() {
+                    let kind = ErrorKind::IncorrectNumberOfParameters {
+                        expected: stmt.params().len(),
+                        actual: params.len(),
+                    };
+
+                    return Err(Error::builder(kind).build());
+                }
 
-            let changes = self
-                .perform_io(self.client.0.execute(&stmt, conversion::conv_params(params).as_slice()))
-                .await?;
+                let changes = self
+                    .perform_io(self.client.0.execute(&stmt, conversion::conv_params(params).as_slice()))
+                    .await?;
 
-            Ok(changes)
-        })
+                Ok(changes)
+            },
+        )
         .await
     }
 
     async fn raw_cmd(&self, cmd: &str) -> crate::Result<()> {
-        metrics::query("postgres.raw_cmd", cmd, &[], move || async move {
+        metrics::query("postgres.raw_cmd", self.db_system_name, cmd, &[], move || async move {
             self.perform_io(self.client.0.simple_query(cmd)).await?;
             Ok(())
         })