matrix/authentication-service
1
0
Fork 0
mirror of https://github.com/matrix-org/matrix-authentication-service.git synced 2025-07-29 22:01:14 +03:00
Code Activity

storage: trace storage operations better

Browse Source
This commit is contained in:
Quentin Gliech
2023-01-02 16:46:34 +01:00
parent 13a9d03647
commit 4790897892
7 changed files with 176 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
crates/storage/src/upstream_oauth2/link.rs
View File

@ -17,12 +17,12 @@ use chrono::{DateTime, Utc};
use mas_data_model::{UpstreamOAuthLink, UpstreamOAuthProvider, User};
use rand::RngCore;
use sqlx::{PgConnection, QueryBuilder};
use tracing::{info_span, Instrument};
use ulid::Ulid;
use uuid::Uuid;
use crate::{
pagination::{process_page, Page, QueryBuilderExt},
tracing::ExecuteExt,
Clock, DatabaseError, LookupResultExt,
};
@ -103,8 +103,12 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
type Error = DatabaseError;
#[tracing::instrument(
name = "db.upstream_oauth_link.lookup",
skip_all,
fields(upstream_oauth_link.id = %id),
fields(
db.statement,
upstream_oauth_link.id = %id,
),
err,
)]
async fn lookup(&mut self, id: Ulid) -> Result<Option<UpstreamOAuthLink>, Self::Error> {
@ -122,6 +126,7 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
"#,
Uuid::from(id),
)
.traced()
.fetch_one(&mut *self.conn)
.await
.to_option()?
@ -131,8 +136,10 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
}
#[tracing::instrument(
name = "db.upstream_oauth_link.find_by_subject",
skip_all,
fields(
db.statement,
upstream_oauth_link.subject = subject,
%upstream_oauth_provider.id,
%upstream_oauth_provider.issuer,
@ -161,6 +168,7 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
Uuid::from(upstream_oauth_provider.id),
subject,
)
.traced()
.fetch_one(&mut *self.conn)
.await
.to_option()?
@ -170,8 +178,10 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
}
#[tracing::instrument(
name = "db.upstream_oauth_link.add",
skip_all,
fields(
db.statement,
upstream_oauth_link.id,
upstream_oauth_link.subject = subject,
%upstream_oauth_provider.id,
@ -206,6 +216,7 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
&subject,
created_at,
)
.traced()
.execute(&mut *self.conn)
.await?;
@ -219,8 +230,10 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
}
#[tracing::instrument(
name = "db.upstream_oauth_link.associate_to_user",
skip_all,
fields(
db.statement,
%upstream_oauth_link.id,
%upstream_oauth_link.subject,
%user.id,
@ -242,6 +255,7 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
Uuid::from(user.id),
Uuid::from(upstream_oauth_link.id),
)
.traced()
.execute(&mut *self.conn)
.await?;
@ -249,8 +263,13 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
}
#[tracing::instrument(
name = "db.upstream_oauth_link.list_paginated",
skip_all,
fields(%user.id, %user.username),
fields(
db.statement,
%user.id,
%user.username,
),
err
)]
async fn list_paginated(
@ -278,14 +297,10 @@ impl<'c> UpstreamOAuthLinkRepository for PgUpstreamOAuthLinkRepository<'c> {
.push_bind(Uuid::from(user.id))
.generate_pagination("upstream_oauth_link_id", before, after, first, last)?;
let span = info_span!(
"Fetch paginated upstream OAuth 2.0 user links",
db.statement = query.sql()
);
let page: Vec<LinkLookup> = query
.build_query_as()
.traced()
.fetch_all(&mut *self.conn)
.instrument(span)
.await?;
let (has_previous_page, has_next_page, edges) = process_page(page, first, last)?;

37
crates/storage/src/upstream_oauth2/provider.rs
View File

@ -19,12 +19,12 @@ use mas_iana::{jose::JsonWebSignatureAlg, oauth::OAuthClientAuthenticationMethod
use oauth2_types::scope::Scope;
use rand::RngCore;
use sqlx::{PgConnection, QueryBuilder};
use tracing::{info_span, Instrument};
use ulid::Ulid;
use uuid::Uuid;
use crate::{
pagination::{process_page, Page, QueryBuilderExt},
tracing::ExecuteExt,
Clock, DatabaseError, DatabaseInconsistencyError, LookupResultExt,
};
@ -129,8 +129,12 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
type Error = DatabaseError;
#[tracing::instrument(
name = "db.upstream_oauth_provider.lookup",
skip_all,
fields(upstream_oauth_provider.id = %id),
fields(
db.statement,
upstream_oauth_provider.id = %id,
),
err,
)]
async fn lookup(&mut self, id: Ulid) -> Result<Option<UpstreamOAuthProvider>, Self::Error> {
@ -151,6 +155,7 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
"#,
Uuid::from(id),
)
.traced()
.fetch_one(&mut *self.conn)
.await
.to_option()?;
@ -164,8 +169,10 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
}
#[tracing::instrument(
name = "db.upstream_oauth_provider.add",
skip_all,
fields(
db.statement,
upstream_oauth_provider.id,
upstream_oauth_provider.issuer = %issuer,
upstream_oauth_provider.client_id = %client_id,
@ -210,6 +217,7 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
encrypted_client_secret.as_deref(),
created_at,
)
.traced()
.execute(&mut *self.conn)
.await?;
@ -225,6 +233,14 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
})
}
#[tracing::instrument(
name = "db.upstream_oauth_provider.list_paginated",
skip_all,
fields(
db.statement,
),
err,
)]
async fn list_paginated(
&mut self,
before: Option<Ulid>,
@ -250,14 +266,10 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
query.generate_pagination("upstream_oauth_provider_id", before, after, first, last)?;
let span = info_span!(
"Fetch paginated upstream OAuth 2.0 providers",
db.statement = query.sql()
);
let page: Vec<ProviderLookup> = query
.build_query_as()
.traced()
.fetch_all(&mut *self.conn)
.instrument(span)
.await?;
let (has_previous_page, has_next_page, edges) = process_page(page, first, last)?;
@ -269,7 +281,15 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
edges: edges?,
})
}
#[tracing::instrument(skip_all, err)]
#[tracing::instrument(
name = "db.upstream_oauth_provider.all",
skip_all,
fields(
db.statement,
),
err,
)]
async fn all(&mut self) -> Result<Vec<UpstreamOAuthProvider>, Self::Error> {
let res = sqlx::query_as!(
ProviderLookup,
@ -286,6 +306,7 @@ impl<'c> UpstreamOAuthProviderRepository for PgUpstreamOAuthProviderRepository<'
FROM upstream_oauth_providers
"#,
)
.traced()
.fetch_all(&mut *self.conn)
.await?;

18
crates/storage/src/upstream_oauth2/session.rs
View File

@ -20,7 +20,7 @@ use sqlx::PgConnection;
use ulid::Ulid;
use uuid::Uuid;
use crate::{Clock, DatabaseError, LookupResultExt};
use crate::{tracing::ExecuteExt, Clock, DatabaseError, LookupResultExt};
#[async_trait]
pub trait UpstreamOAuthSessionRepository: Send + Sync {
@ -88,8 +88,12 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
type Error = DatabaseError;
#[tracing::instrument(
name = "db.upstream_oauth_authorization_session.lookup",
skip_all,
fields(upstream_oauth_provider.id = %id),
fields(
db.statement,
upstream_oauth_provider.id = %id,
),
err,
)]
async fn lookup(
@ -115,6 +119,7 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
"#,
Uuid::from(id),
)
.traced()
.fetch_one(&mut *self.conn)
.await
.to_option()?;
@ -138,8 +143,10 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
}
#[tracing::instrument(
name = "db.upstream_oauth_authorization_session.add",
skip_all,
fields(
db.statement,
%upstream_oauth_provider.id,
%upstream_oauth_provider.issuer,
%upstream_oauth_provider.client_id,
@ -184,6 +191,7 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
nonce,
created_at,
)
.traced()
.execute(&mut *self.conn)
.await?;
@ -202,8 +210,10 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
}
#[tracing::instrument(
name = "db.upstream_oauth_authorization_session.complete_with_link",
skip_all,
fields(
db.statement,
%upstream_oauth_authorization_session.id,
%upstream_oauth_link.id,
),
@ -230,6 +240,7 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
id_token,
Uuid::from(upstream_oauth_authorization_session.id),
)
.traced()
.execute(&mut *self.conn)
.await?;
@ -242,8 +253,10 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
/// Mark a session as consumed
#[tracing::instrument(
name = "db.upstream_oauth_authorization_session.consume",
skip_all,
fields(
db.statement,
%upstream_oauth_authorization_session.id,
),
err,
@ -263,6 +276,7 @@ impl<'c> UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'c>
consumed_at,
Uuid::from(upstream_oauth_authorization_session.id),
)
.traced()
.execute(&mut *self.conn)
.await?;