Add SQLx operation tracing (#5223)

* wip: vendor sqlx-tracing

* (compiles) standardize pg types used

* more standardization

* general log message improvements

* wip: improve sqlx-tracing architecture

* unify sqlx::Executor type

* wip: try fix sqlx tracing

* wip: sqlx-tracing compiles

* so close

* it compiles

* fix ci

packages/sqlx-tracing/tests/api.rs

@@ -0,0 +1,87 @@
+//! Test that valid uses of the API compile.
+#![expect(dead_code, reason = "only here to check that the code compiles")]
+
+use sqlx::Postgres;
+
+async fn a(db: sqlx_tracing::Pool<Postgres>) {
+    let _conn: sqlx_tracing::PoolConnection<Postgres> =
+        db.acquire().await.unwrap();
+}
+
+async fn b<'a, E>(exec: E)
+where
+    E: sqlx_tracing::Acquire<'a, Database = Postgres>,
+{
+    let mut conn: sqlx_tracing::AnyConnection<Postgres> =
+        exec.acquire().await.unwrap();
+    // sqlx::query("SELECT 1").execute(&mut conn).await.unwrap();
+    sqlx::query("SELECT 1").execute(&mut conn).await.unwrap();
+}
+
+async fn c<'a, E>(exec: E)
+where
+    E: sqlx_tracing::Executor<'a, Database = Postgres>,
+{
+    sqlx::query("SELECT 1").execute(exec).await.unwrap();
+}
+
+pub async fn list_many<'a, E>(exec: E)
+where
+    E: sqlx::Executor<'a, Database = Postgres>,
+{
+    sqlx::query(
+        "
+        SELECT
+            id, enum_id, value, ordering,
+            metadata, created
+        FROM loader_field_enum_values
+        WHERE enum_id = ANY($1)
+        ORDER BY enum_id, ordering, created DESC
+        ",
+    )
+    .fetch_all(exec)
+    .await
+    .unwrap();
+}
+
+async fn insert_sqlx(transaction: &mut sqlx::Transaction<'_, Postgres>) {
+    get_id_sqlx(&mut *transaction).await;
+}
+
+async fn insert<'t, 'c>(
+    transaction: &'t mut sqlx_tracing::Transaction<'c, Postgres>,
+) {
+    get_id(&mut *transaction).await;
+    get_id(&mut *transaction).await;
+
+    sqlx::query("SELECT 1")
+        .execute(&mut *transaction)
+        .await
+        .unwrap();
+    sqlx::query("SELECT 1")
+        .execute(&mut *transaction)
+        .await
+        .unwrap();
+}
+
+async fn get_id_sqlx<'a, E>(_executor: E)
+where
+    E: sqlx::Acquire<'a, Database = sqlx::Postgres>,
+{
+}
+
+async fn get_id<'a, E>(_executor: E)
+where
+    E: sqlx_tracing::Acquire<'a, Database = sqlx::Postgres>,
+{
+}
+
+async fn d<'a, E>(exec: E)
+where
+    E: sqlx_tracing::Acquire<'a, Database = Postgres>,
+{
+    let mut exec = exec.acquire().await.unwrap();
+
+    sqlx::query("SELECT 1").fetch_one(&mut exec).await.unwrap();
+    sqlx::query("SELECT 1").fetch_one(&mut exec).await.unwrap();
+}

packages/sqlx-tracing/tests/common.rs

@@ -0,0 +1,47 @@
+use std::time::Duration;
+
+use opentelemetry::trace::{FutureExt, TraceContextExt, Tracer};
+
+pub async fn should_trace<'c, DB, E>(
+    name: &'static str,
+    system: &'static str,
+    observability: &opentelemetry_testing::ObservabilityContainer,
+    provider: &opentelemetry_testing::OpenTelemetryProvider,
+    executor: E,
+) where
+    DB: sqlx::Database,
+    E: sqlx::Executor<'c, Database = DB>,
+    for<'q> DB::Arguments<'q>: 'q + sqlx::IntoArguments<'q, DB>,
+    (i32,): Send + Unpin + for<'r> sqlx::FromRow<'r, DB::Row>,
+{
+    let scope = format!("should_{name}_{system}");
+    let tracer = opentelemetry::global::tracer(scope.clone());
+    let span = tracer.span_builder(name).start(&tracer);
+    let ctx = opentelemetry::Context::new().with_span(span);
+
+    let result: Option<i32> = sqlx::query_scalar("select 1")
+        .fetch_optional(executor)
+        .with_context(ctx)
+        .await
+        .unwrap();
+
+    assert_eq!(result, Some(1));
+
+    provider.flush();
+
+    tokio::time::sleep(Duration::from_secs(1)).await;
+
+    let traces = observability.json_traces();
+    let scope_span = traces.find_scope_span(&scope).unwrap();
+    let entry = scope_span.first_span().unwrap();
+    assert_eq!(entry.name, name);
+    let next = traces
+        .find_child(&entry.span_id, "sqlx.fetch_optional")
+        .unwrap();
+    assert_eq!(next.string_attribute("db.system.name").unwrap(), system);
+    assert_eq!(next.string_attribute("db.query.text").unwrap(), "select 1");
+    assert_eq!(
+        next.int_attribute("db.response.returned_rows").unwrap(),
+        "1"
+    );
+}

packages/sqlx-tracing/tests/postgres.rs

@@ -0,0 +1,80 @@
+#![cfg(feature = "postgres")]
+
+use std::time::Duration;
+
+use sqlx::Postgres;
+use testcontainers::{
+    GenericImage, ImageExt,
+    core::{ContainerPort, WaitFor},
+    runners::AsyncRunner,
+};
+
+mod common;
+
+#[derive(Debug)]
+struct PostgresContainer {
+    container: testcontainers::ContainerAsync<testcontainers::GenericImage>,
+}
+
+impl PostgresContainer {
+    async fn create() -> Self {
+        let container = GenericImage::new("postgres", "15-alpine")
+            .with_wait_for(WaitFor::message_on_stderr(
+                "database system is ready to accept connections",
+            ))
+            .with_exposed_port(ContainerPort::Tcp(5432))
+            .with_env_var("POSTGRES_USER", "postgres")
+            .with_env_var("POSTGRES_DB", "postgres")
+            .with_env_var("POSTGRES_HOST_AUTH_METHOD", "trust")
+            .with_startup_timeout(Duration::from_secs(60))
+            .start()
+            .await
+            .expect("starting a postgres database");
+
+        Self { container }
+    }
+
+    async fn client(&self) -> sqlx_tracing::Pool<Postgres> {
+        let port = self.container.get_host_port_ipv4(5432).await.unwrap();
+        let url = format!("postgres://postgres@localhost:{port}/postgres");
+        sqlx::PgPool::connect(&url)
+            .await
+            .map(sqlx_tracing::Pool::from)
+            .unwrap()
+    }
+}
+
+#[tokio::test]
+async fn execute() {
+    let observability = opentelemetry_testing::ObservabilityContainer::create().await;
+    let provider = observability.install().await;
+
+    let container = PostgresContainer::create().await;
+    let pool = container.client().await;
+
+    common::should_trace("trace_pool", "postgresql", &observability, &provider, &pool).await;
+
+    {
+        let mut conn = pool.acquire().await.unwrap();
+        common::should_trace(
+            "trace_conn",
+            "postgresql",
+            &observability,
+            &provider,
+            &mut conn,
+        )
+        .await;
+    }
+
+    {
+        let mut tx: sqlx_tracing::Transaction<'_, Postgres> = pool.begin().await.unwrap();
+        common::should_trace(
+            "trace_tx",
+            "postgresql",
+            &observability,
+            &provider,
+            &mut tx.executor(),
+        )
+        .await;
+    }
+}

packages/sqlx-tracing/tests/sqlite.rs

@@ -0,0 +1,33 @@
+#![cfg(feature = "sqlite")]
+
+use sqlx::Sqlite;
+
+mod common;
+
+#[tokio::test]
+async fn execute() {
+    let observability = opentelemetry_testing::ObservabilityContainer::create().await;
+    let provider = observability.install().await;
+
+    let pool = sqlx::SqlitePool::connect(":memory:").await.unwrap();
+    let pool = sqlx_tracing::Pool::from(pool);
+
+    common::should_trace("trace_pool", "sqlite", &observability, &provider, &pool).await;
+
+    {
+        let mut conn = pool.acquire().await.unwrap();
+        common::should_trace("trace_conn", "sqlite", &observability, &provider, &mut conn).await;
+    }
+
+    {
+        let mut tx: sqlx_tracing::Transaction<'_, Sqlite> = pool.begin().await.unwrap();
+        common::should_trace(
+            "trace_tx",
+            "sqlite",
+            &observability,
+            &provider,
+            &mut tx.executor(),
+        )
+        .await;
+    }
+}