Skip to content

Commit 31ccd54

Browse files
committed
Properly record the queries in two spans
1 parent a73cb1c commit 31ccd54

File tree

1 file changed

+14
-3
lines changed

1 file changed

+14
-3
lines changed

crates/storage-pg/src/app_session.rs

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,13 @@ use mas_storage::{
1717
oauth2::OAuth2SessionFilter,
1818
};
1919
use oauth2_types::scope::{Scope, ScopeToken};
20+
use opentelemetry_semantic_conventions::trace::DB_QUERY_TEXT;
2021
use sea_query::{
2122
Alias, ColumnRef, CommonTableExpression, Expr, PostgresQueryBuilder, Query, UnionType,
2223
};
2324
use sea_query_binder::SqlxBinder;
2425
use sqlx::PgConnection;
26+
use tracing::Instrument;
2527
use ulid::Ulid;
2628
use uuid::Uuid;
2729

@@ -478,7 +480,10 @@ impl AppSessionRepository for PgAppSessionRepository<'_> {
478480
device: &Device,
479481
) -> Result<(), Self::Error> {
480482
// TODO need to invoke this from all the oauth2 login sites
481-
// TODO CREATE A SECOND SPAN FOR THE SECOND QUERY
483+
let span = tracing::info_span!(
484+
"db.app_session.finish_sessions_to_replace_device.compat_sessions",
485+
{ DB_QUERY_TEXT } = tracing::field::Empty,
486+
);
482487
let finished_at = clock.now();
483488
sqlx::query!(
484489
"
@@ -488,11 +493,16 @@ impl AppSessionRepository for PgAppSessionRepository<'_> {
488493
device.as_str(),
489494
finished_at
490495
)
491-
.traced()
496+
.record(&span)
492497
.execute(&mut *self.conn)
498+
.instrument(span)
493499
.await?;
494500

495501
if let Ok(device_as_scope_token) = device.to_scope_token() {
502+
let span = tracing::info_span!(
503+
"db.app_session.finish_sessions_to_replace_device.oauth2_sessions",
504+
{ DB_QUERY_TEXT } = tracing::field::Empty,
505+
);
496506
sqlx::query!(
497507
"
498508
UPDATE oauth2_sessions SET finished_at = $3 WHERE user_id = $1 AND $2 = ANY(scope_list) AND finished_at IS NULL
@@ -501,8 +511,9 @@ impl AppSessionRepository for PgAppSessionRepository<'_> {
501511
device_as_scope_token.as_str(),
502512
finished_at
503513
)
504-
.traced()
514+
.record(&span)
505515
.execute(&mut *self.conn)
516+
.instrument(span)
506517
.await?;
507518
}
508519

0 commit comments

Comments
 (0)