From ae2fe1cb3f61b8e0633b4692f74b2cc8a7a6fe93 Mon Sep 17 00:00:00 2001
From: Tonkku <contact@tonkku.me>
Date: Tue, 8 Apr 2025 16:54:35 +0000
Subject: [PATCH] Separate spans

---
 crates/storage-pg/src/upstream_oauth2/link.rs | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/crates/storage-pg/src/upstream_oauth2/link.rs b/crates/storage-pg/src/upstream_oauth2/link.rs
index ea9cda163..390029f1f 100644
--- a/crates/storage-pg/src/upstream_oauth2/link.rs
+++ b/crates/storage-pg/src/upstream_oauth2/link.rs
@@ -11,10 +11,12 @@ use mas_storage::{
     Clock, Page, Pagination,
     upstream_oauth2::{UpstreamOAuthLinkFilter, UpstreamOAuthLinkRepository},
 };
+use opentelemetry_semantic_conventions::trace::DB_QUERY_TEXT;
 use rand::RngCore;
 use sea_query::{Expr, PostgresQueryBuilder, Query, enum_def};
 use sea_query_binder::SqlxBinder;
 use sqlx::PgConnection;
+use tracing::Instrument;
 use ulid::Ulid;
 use uuid::Uuid;
 
@@ -393,6 +395,10 @@ impl UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'_> {
     ) -> Result<(), Self::Error> {
         // Unlink the authorization sessions first, as they have a foreign key
         // constraint on the links.
+        let span = tracing::info_span!(
+            "db.upstream_oauth_link.remove.unlink",
+            { DB_QUERY_TEXT } = tracing::field::Empty
+        );
         sqlx::query!(
             r#"
                 UPDATE upstream_oauth_authorization_sessions SET
@@ -403,11 +409,16 @@ impl UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'_> {
             Uuid::from(upstream_oauth_link.id),
             clock.now()
         )
-        .traced()
+        .record(&span)
         .execute(&mut *self.conn)
+        .instrument(span)
         .await?;
 
         // Then delete the link itself
+        let span = tracing::info_span!(
+            "db.upstream_oauth_link.remove.delete",
+            { DB_QUERY_TEXT } = tracing::field::Empty
+        );
         let res = sqlx::query!(
             r#"
                 DELETE FROM upstream_oauth_links
@@ -415,8 +426,9 @@ impl UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'_> {
             "#,
             Uuid::from(upstream_oauth_link.id),
         )
-        .traced()
+        .record(&span)
         .execute(&mut *self.conn)
+        .instrument(span)
         .await?;
 
         DatabaseError::ensure_affected_rows(&res, 1)?;