Better DB operations tracing

crates/storage/sqlx-data.json

@@ -291,19 +291,6 @@
     },
     "query": "\n                INSERT INTO oauth2_client_redirect_uris\n                    (oauth2_client_redirect_uri_id, oauth2_client_id, redirect_uri)\n                VALUES ($1, $2, $3)\n            "
   },
-  "1a10e6189300563e79684eb7ccc6c29b0418aadfdeea6f8bc5a700a411409c73": {
-    "describe": {
-      "columns": [],
-      "nullable": [],
-      "parameters": {
-        "Left": [
-          "Text",
-          "Timestamptz"
-        ]
-      }
-    },
-    "query": "\n            UPDATE compat_sessions cs\n            SET finished_at = $2\n            FROM compat_access_tokens ca\n            WHERE ca.access_token = $1\n              AND ca.compat_session_id = cs.compat_session_id\n              AND cs.finished_at IS NULL\n        "
-  },
   "1eb6d13e75d8f526c2785749a020731c18012f03e07995213acd38ab560ce497": {
     "describe": {
       "columns": [],
@@ -1023,6 +1010,27 @@
     },
     "query": "\n            SELECT scope_token\n            FROM oauth2_consents\n            WHERE user_id = $1 AND oauth2_client_id = $2\n        "
   },
+  "559a486756d08d101eb7188ef6637b9d24c024d056795b8121f7f04a7f9db6a3": {
+    "describe": {
+      "columns": [
+        {
+          "name": "compat_session_id",
+          "ordinal": 0,
+          "type_info": "Uuid"
+        }
+      ],
+      "nullable": [
+        false
+      ],
+      "parameters": {
+        "Left": [
+          "Text",
+          "Timestamptz"
+        ]
+      }
+    },
+    "query": "\n            UPDATE compat_sessions cs\n            SET finished_at = $2\n            FROM compat_access_tokens ca\n            WHERE ca.access_token = $1\n              AND ca.compat_session_id = cs.compat_session_id\n              AND cs.finished_at IS NULL\n            RETURNING cs.compat_session_id\n        "
+  },
   "5b5d5c82da37c6f2d8affacfb02119965c04d1f2a9cc53dbf5bd4c12584969a0": {
     "describe": {
       "columns": [],

crates/storage/src/compat.rs

@@ -300,7 +300,16 @@ pub async fn lookup_active_compat_refresh_token(
     Ok((refresh_token, access_token, session))
 }
 
-#[tracing::instrument(skip(conn, password), err)]
+#[tracing::instrument(
+    skip_all,
+    fields(
+        user.username = username,
+        user.id,
+        compat_session.id,
+        compat_session.device.id = device.as_str(),
+    ),
+    err(Display),
+)]
 pub async fn compat_login(
     conn: impl Acquire<'_, Database = Postgres>,
     username: &str,
@@ -311,6 +320,7 @@ pub async fn compat_login(
 
     // First, lookup the user
     let user = lookup_user_by_username(&mut txn, username).await?;
+    tracing::Span::current().record("user.id", tracing::field::display(user.data));
 
     // Now, fetch the hashed password from the user associated with that session
     let hashed_password: String = sqlx::query_scalar!(
@@ -340,6 +350,8 @@ pub async fn compat_login(
 
     let created_at = Utc::now();
     let id = Ulid::from_datetime(created_at.into());
+    tracing::Span::current().record("compat_session.id", tracing::field::display(id));
+
     sqlx::query!(
         r#"
             INSERT INTO compat_sessions
@@ -368,7 +380,16 @@ pub async fn compat_login(
     Ok(session)
 }
 
-#[tracing::instrument(skip(executor, token), err)]
+#[tracing::instrument(
+    skip_all,
+    fields(
+        compat_session.id = %session.data,
+        compat_session.device.id = session.device.as_str(),
+        compat_access_token.id,
+        user.id = %session.user.data,
+    ),
+    err(Display),
+)]
 pub async fn add_compat_access_token(
     executor: impl PgExecutor<'_>,
     session: &CompatSession<PostgresqlBackend>,
@@ -377,6 +398,8 @@ pub async fn add_compat_access_token(
 ) -> Result<CompatAccessToken<PostgresqlBackend>, anyhow::Error> {
     let created_at = Utc::now();
     let id = Ulid::from_datetime(created_at.into());
+    tracing::Span::current().record("compat_access_token.id", tracing::field::display(id));
+
     let expires_at = expires_after.map(|expires_after| created_at + expires_after);
 
     sqlx::query!(
@@ -404,10 +427,17 @@ pub async fn add_compat_access_token(
     })
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        compat_access_token.id = %access_token.data,
+    ),
+    err(Display),
+)]
 pub async fn expire_compat_access_token(
     executor: impl PgExecutor<'_>,
     access_token: CompatAccessToken<PostgresqlBackend>,
-) -> anyhow::Result<()> {
+) -> Result<(), anyhow::Error> {
     let expires_at = Utc::now();
     let res = sqlx::query!(
         r#"
@@ -431,6 +461,17 @@ pub async fn expire_compat_access_token(
     }
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        compat_session.id = %session.data,
+        compat_session.device.id = session.device.as_str(),
+        compat_access_token.id = %access_token.data,
+        compat_refresh_token.id,
+        user.id = %session.user.data,
+    ),
+    err(Display),
+)]
 pub async fn add_compat_refresh_token(
     executor: impl PgExecutor<'_>,
     session: &CompatSession<PostgresqlBackend>,
@@ -439,6 +480,8 @@ pub async fn add_compat_refresh_token(
 ) -> Result<CompatRefreshToken<PostgresqlBackend>, anyhow::Error> {
     let created_at = Utc::now();
     let id = Ulid::from_datetime(created_at.into());
+    tracing::Span::current().record("compat_refresh_token.id", tracing::field::display(id));
+
     sqlx::query!(
         r#"
             INSERT INTO compat_refresh_tokens
@@ -464,14 +507,18 @@ pub async fn add_compat_refresh_token(
     })
 }
 
-#[tracing::instrument(skip_all, err)]
+#[tracing::instrument(
+    skip_all,
+    fields(compat_session.id),
+    err(Display),
+)]
 pub async fn compat_logout(
     executor: impl PgExecutor<'_>,
     token: &str,
 ) -> Result<(), anyhow::Error> {
     let finished_at = Utc::now();
     // TODO: this does not check for token expiration
-    let res = sqlx::query!(
+    let compat_session_id = sqlx::query_scalar!(
         r#"
             UPDATE compat_sessions cs
             SET finished_at = $2
@@ -479,25 +526,34 @@ pub async fn compat_logout(
             WHERE ca.access_token = $1
               AND ca.compat_session_id = cs.compat_session_id
               AND cs.finished_at IS NULL
+            RETURNING cs.compat_session_id
         "#,
         token,
         finished_at,
     )
-    .execute(executor)
+    .fetch_one(executor)
     .await
     .context("could not update compat access token")?;
 
-    match res.rows_affected() {
+    tracing::Span::current().record(
-        1 => Ok(()),
+        "compat_session.id",
-        0 => anyhow::bail!("no row affected"),
+        tracing::field::display(compat_session_id),
-        _ => anyhow::bail!("too many row affected"),
+    );
-    }
+
+    Ok(())
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        compat_refresh_token.id = %refresh_token.data,
+    ),
+    err(Display),
+)]
 pub async fn consume_compat_refresh_token(
     executor: impl PgExecutor<'_>,
     refresh_token: CompatRefreshToken<PostgresqlBackend>,
-) -> anyhow::Result<()> {
+) -> Result<(), anyhow::Error> {
     let consumed_at = Utc::now();
     let res = sqlx::query!(
         r#"
@@ -521,13 +577,23 @@ pub async fn consume_compat_refresh_token(
     }
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        compat_sso_login.id,
+        compat_sso_login.redirect_uri = %redirect_uri,
+    ),
+    err(Display),
+)]
 pub async fn insert_compat_sso_login(
     executor: impl PgExecutor<'_>,
     login_token: String,
     redirect_uri: Url,
-) -> anyhow::Result<CompatSsoLogin<PostgresqlBackend>> {
+) -> Result<CompatSsoLogin<PostgresqlBackend>, anyhow::Error> {
     let created_at = Utc::now();
     let id = Ulid::from_datetime(created_at.into());
+    tracing::Span::current().record("compat_sso_login.id", tracing::field::display(id));
+
     sqlx::query!(
         r#"
             INSERT INTO compat_sso_logins
@@ -675,8 +741,13 @@ impl CompatSsoLoginLookupError {
     }
 }
 
-#[allow(clippy::too_many_lines)]
+#[tracing::instrument(
-#[tracing::instrument(skip(executor), err)]
+    skip_all,
+    fields(
+        compat_sso_login.id = %id,
+    ),
+    err,
+)]
 pub async fn get_compat_sso_login_by_id(
     executor: impl PgExecutor<'_>,
     id: Ulid,
@@ -719,8 +790,7 @@ pub async fn get_compat_sso_login_by_id(
     Ok(res.try_into()?)
 }
 
-#[allow(clippy::too_many_lines)]
+#[tracing::instrument(skip_all, err)]
-#[tracing::instrument(skip(executor), err)]
 pub async fn get_compat_sso_login_by_token(
     executor: impl PgExecutor<'_>,
     token: &str,
@@ -763,12 +833,23 @@ pub async fn get_compat_sso_login_by_token(
     Ok(res.try_into()?)
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        user.id = %user.data,
+        compat_sso_login.id = %login.data,
+        compat_sso_login.redirect_uri = %login.redirect_uri,
+        compat_session.id,
+        compat_session.device.id = device.as_str(),
+    ),
+    err(Display),
+)]
 pub async fn fullfill_compat_sso_login(
     conn: impl Acquire<'_, Database = Postgres>,
     user: User<PostgresqlBackend>,
     mut login: CompatSsoLogin<PostgresqlBackend>,
     device: Device,
-) -> anyhow::Result<CompatSsoLogin<PostgresqlBackend>> {
+) -> Result<CompatSsoLogin<PostgresqlBackend>, anyhow::Error> {
     if !matches!(login.state, CompatSsoLoginState::Pending) {
         bail!("sso login in wrong state");
     };
@@ -777,6 +858,8 @@ pub async fn fullfill_compat_sso_login(
 
     let created_at = Utc::now();
     let id = Ulid::from_datetime(created_at.into());
+    tracing::Span::current().record("user.id", tracing::field::display(user.data));
+
     sqlx::query!(
         r#"
             INSERT INTO compat_sessions (compat_session_id, user_id, device_id, created_at)
@@ -831,10 +914,18 @@ pub async fn fullfill_compat_sso_login(
     Ok(login)
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        compat_sso_login.id = %login.data,
+        compat_sso_login.redirect_uri = %login.redirect_uri,
+    ),
+    err(Display),
+)]
 pub async fn mark_compat_sso_login_as_exchanged(
     executor: impl PgExecutor<'_>,
     mut login: CompatSsoLogin<PostgresqlBackend>,
-) -> anyhow::Result<CompatSsoLogin<PostgresqlBackend>> {
+) -> Result<CompatSsoLogin<PostgresqlBackend>, anyhow::Error> {
     let (fulfilled_at, session) = match login.state {
         CompatSsoLoginState::Fulfilled {
             fulfilled_at,

crates/storage/src/oauth2/access_token.rs

@@ -38,7 +38,7 @@ pub async fn add_access_token(
     session: &Session<PostgresqlBackend>,
     access_token: String,
     expires_after: Duration,
-) -> anyhow::Result<AccessToken<PostgresqlBackend>> {
+) -> Result<AccessToken<PostgresqlBackend>, anyhow::Error> {
     let created_at = Utc::now();
     let expires_at = created_at + expires_after;
     let id = Ulid::from_datetime(created_at.into());

crates/storage/src/oauth2/authorization_grant.rs

@@ -54,7 +54,7 @@ pub async fn new_authorization_grant(
     response_mode: ResponseMode,
     response_type_id_token: bool,
     requires_consent: bool,
-) -> anyhow::Result<AuthorizationGrant<PostgresqlBackend>> {
+) -> Result<AuthorizationGrant<PostgresqlBackend>, anyhow::Error> {
     let code_challenge = code
         .as_ref()
         .and_then(|c| c.pkce.as_ref())
@@ -359,7 +359,7 @@ impl GrantLookup {
 pub async fn get_grant_by_id(
     conn: &mut PgConnection,
     id: Ulid,
-) -> anyhow::Result<AuthorizationGrant<PostgresqlBackend>> {
+) -> Result<AuthorizationGrant<PostgresqlBackend>, anyhow::Error> {
     // TODO: handle "not found" cases
     let res = sqlx::query_as!(
         GrantLookup,
@@ -427,7 +427,7 @@ pub async fn get_grant_by_id(
 pub async fn lookup_grant_by_code(
     conn: &mut PgConnection,
     code: &str,
-) -> anyhow::Result<AuthorizationGrant<PostgresqlBackend>> {
+) -> Result<AuthorizationGrant<PostgresqlBackend>, anyhow::Error> {
     // TODO: handle "not found" cases
     let res = sqlx::query_as!(
         GrantLookup,
@@ -506,7 +506,7 @@ pub async fn derive_session(
     executor: impl PgExecutor<'_>,
     grant: &AuthorizationGrant<PostgresqlBackend>,
     browser_session: BrowserSession<PostgresqlBackend>,
-) -> anyhow::Result<Session<PostgresqlBackend>> {
+) -> Result<Session<PostgresqlBackend>, anyhow::Error> {
     let created_at = Utc::now();
     let id = Ulid::from_datetime(created_at.into());
     tracing::Span::current().record("session.id", tracing::field::display(id));
@@ -558,7 +558,7 @@ pub async fn fulfill_grant(
     executor: impl PgExecutor<'_>,
     mut grant: AuthorizationGrant<PostgresqlBackend>,
     session: Session<PostgresqlBackend>,
-) -> anyhow::Result<AuthorizationGrant<PostgresqlBackend>> {
+) -> Result<AuthorizationGrant<PostgresqlBackend>, anyhow::Error> {
     let fulfilled_at = sqlx::query_scalar!(
         r#"
             UPDATE oauth2_authorization_grants AS og
@@ -624,7 +624,7 @@ pub async fn give_consent_to_grant(
 pub async fn exchange_grant(
     executor: impl PgExecutor<'_>,
     mut grant: AuthorizationGrant<PostgresqlBackend>,
-) -> anyhow::Result<AuthorizationGrant<PostgresqlBackend>> {
+) -> Result<AuthorizationGrant<PostgresqlBackend>, anyhow::Error> {
     let exchanged_at = Utc::now();
     sqlx::query!(
         r#"

crates/storage/src/oauth2/client.rs

@@ -248,6 +248,11 @@ impl TryInto<Client<PostgresqlBackend>> for OAuth2ClientLookup {
     }
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(client.id = %id),
+    err,
+)]
 pub async fn lookup_client(
     executor: impl PgExecutor<'_>,
     id: Ulid,
@@ -291,6 +296,11 @@ pub async fn lookup_client(
     Ok(client)
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(client.id = client_id),
+    err,
+)]
 pub async fn lookup_client_by_client_id(
     executor: impl PgExecutor<'_>,
     client_id: &str,
@@ -299,6 +309,11 @@ pub async fn lookup_client_by_client_id(
     lookup_client(executor, id).await
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(client.id = %client_id, client.name = client_name),
+    err,
+)]
 #[allow(clippy::too_many_arguments)]
 pub async fn insert_client(
     conn: &mut PgConnection,
@@ -403,7 +418,7 @@ pub async fn insert_client_from_config(
     jwks: Option<&PublicJsonWebKeySet>,
     jwks_uri: Option<&Url>,
     redirect_uris: &[Url],
-) -> anyhow::Result<()> {
+) -> Result<(), anyhow::Error> {
     let jwks = jwks.map(serde_json::to_value).transpose()?;
     let jwks_uri = jwks_uri.map(Url::as_str);
 
@@ -452,7 +467,7 @@ pub async fn insert_client_from_config(
     Ok(())
 }
 
-pub async fn truncate_clients(executor: impl PgExecutor<'_>) -> anyhow::Result<()> {
+pub async fn truncate_clients(executor: impl PgExecutor<'_>) -> Result<(), anyhow::Error> {
     sqlx::query!("TRUNCATE oauth2_client_redirect_uris, oauth2_clients CASCADE")
         .execute(executor)
         .await?;

crates/storage/src/oauth2/consent.rs

@@ -23,11 +23,19 @@ use uuid::Uuid;
 
 use crate::PostgresqlBackend;
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        user.id = %user.data,
+        client.id = %client.data,
+    ),
+    err(Debug),
+)]
 pub async fn fetch_client_consent(
     executor: impl PgExecutor<'_>,
     user: &User<PostgresqlBackend>,
     client: &Client<PostgresqlBackend>,
-) -> anyhow::Result<Scope> {
+) -> Result<Scope, anyhow::Error> {
     let scope_tokens: Vec<String> = sqlx::query_scalar!(
         r#"
             SELECT scope_token
@@ -48,12 +56,21 @@ pub async fn fetch_client_consent(
     Ok(scope?)
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        user.id = %user.data,
+        client.id = %client.data,
+        scope = scope.to_string(),
+    ),
+    err(Debug),
+)]
 pub async fn insert_client_consent(
     executor: impl PgExecutor<'_>,
     user: &User<PostgresqlBackend>,
     client: &Client<PostgresqlBackend>,
     scope: &Scope,
-) -> anyhow::Result<()> {
+) -> Result<(), anyhow::Error> {
     let now = Utc::now();
     let (tokens, ids): (Vec<String>, Vec<Uuid>) = scope
         .iter()

crates/storage/src/oauth2/mod.rs

@@ -25,10 +25,20 @@ pub mod client;
 pub mod consent;
 pub mod refresh_token;
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        session.id = %session.data,
+        user.id = %session.browser_session.user.data,
+        user_session.id = %session.browser_session.data,
+        client.id = %session.client.data,
+    ),
+    err(Debug),
+)]
 pub async fn end_oauth_session(
     executor: impl PgExecutor<'_>,
     session: Session<PostgresqlBackend>,
-) -> anyhow::Result<()> {
+) -> Result<(), anyhow::Error> {
     let finished_at = Utc::now();
     let res = sqlx::query!(
         r#"

crates/storage/src/oauth2/refresh_token.rs

@@ -25,6 +25,17 @@ use uuid::Uuid;
 use super::client::{lookup_client, ClientFetchError};
 use crate::{DatabaseInconsistencyError, PostgresqlBackend};
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        session.id = %session.data,
+        user.id = %session.browser_session.user.data,
+        user_session.id = %session.browser_session.data,
+        client.id = %session.client.data,
+        refresh_token.id,
+    ),
+    err(Debug),
+)]
 pub async fn add_refresh_token(
     executor: impl PgExecutor<'_>,
     session: &Session<PostgresqlBackend>,
@@ -33,6 +44,7 @@ pub async fn add_refresh_token(
 ) -> anyhow::Result<RefreshToken<PostgresqlBackend>> {
     let created_at = Utc::now();
     let id = Ulid::from_datetime(created_at.into());
+    tracing::Span::current().record("refresh_token.id", tracing::field::display(id));
 
     sqlx::query!(
         r#"
@@ -98,6 +110,7 @@ impl RefreshTokenLookupError {
     }
 }
 
+#[tracing::instrument(skip_all, err)]
 #[allow(clippy::too_many_lines)]
 pub async fn lookup_active_refresh_token(
     conn: &mut PgConnection,
@@ -241,10 +254,17 @@ pub async fn lookup_active_refresh_token(
     Ok((refresh_token, session))
 }
 
+#[tracing::instrument(
+    skip_all,
+    fields(
+        refresh_token.id = %refresh_token.data,
+    ),
+    err(Debug),
+)]
 pub async fn consume_refresh_token(
     executor: impl PgExecutor<'_>,
     refresh_token: &RefreshToken<PostgresqlBackend>,
-) -> anyhow::Result<()> {
+) -> Result<(), anyhow::Error> {
     let consumed_at = Utc::now();
     let res = sqlx::query!(
         r#"