improve perf of setting api roles

This commit is contained in:
Ruben Fiszel
2025-04-12 14:47:48 +00:00
parent 8c6e620f5c
commit fd47cd6016
4 changed files with 92 additions and 43 deletions
@@ -0,0 +1 @@
-- Add down migration script here
@@ -0,0 +1,22 @@
-- Add up migration script here
CREATE OR REPLACE FUNCTION set_session_context(
admin BOOLEAN,
username TEXT,
groups TEXT,
pgroups TEXT,
folders_read TEXT,
folders_write TEXT
) RETURNS void AS $$
BEGIN
IF admin THEN
SET LOCAL ROLE windmill_admin;
ELSE
SET LOCAL ROLE windmill_user;
END IF;
PERFORM set_config('session.user', username, true);
PERFORM set_config('session.groups', groups, true);
PERFORM set_config('session.pgroups', pgroups, true);
PERFORM set_config('session.folders_read', folders_read, true);
PERFORM set_config('session.folders_write', folders_write, true);
END;
$$ LANGUAGE plpgsql;
+64 -41
View File
@@ -72,12 +72,6 @@ impl UserDB {
where
T: Authable,
{
let user = if authed.is_admin() {
"windmill_admin"
} else {
"windmill_user"
};
let (folders_write, folders_read): &(Vec<_>, Vec<_>) =
&authed.folders().into_iter().partition(|x| x.1);
@@ -95,10 +89,6 @@ impl UserDB {
let mut tx = self.db.begin().await?;
sqlx::query(&format!("SET LOCAL ROLE {}", user))
.execute(&mut *tx)
.await?;
if let Some(schema) = PG_SCHEMA.as_ref() {
sqlx::query(&format!("SET LOCAL search_path TO {}", schema))
.execute(&mut *tx)
@@ -106,53 +96,86 @@ impl UserDB {
}
sqlx::query!(
"SELECT set_config('session.user', $1, true)",
authed.username()
)
.fetch_optional(&mut *tx)
.await?;
sqlx::query!(
"SELECT set_config('session.groups', $1, true)",
&authed.groups().join(",")
)
.fetch_optional(&mut *tx)
.await?;
sqlx::query!(
"SELECT set_config('session.pgroups', $1, true)",
&authed
"SELECT set_session_context($1, $2, $3, $4, $5, $6)",
authed.is_admin(),
authed.username(),
authed.groups().join(","),
authed
.groups()
.iter()
.map(|x| format!("g/{}", x))
.collect::<Vec<_>>()
.join(",")
)
.fetch_optional(&mut *tx)
.await?;
sqlx::query!(
"SELECT set_config('session.folders_read', $1, true)",
.join(","),
folders_read
.iter()
.map(|x| x.0.clone())
.collect::<Vec<_>>()
.join(",")
)
.fetch_optional(&mut *tx)
.await?;
sqlx::query!(
"SELECT set_config('session.folders_write', $1, true)",
.join(","),
folders_write
.iter()
.map(|x| x.0.clone())
.collect::<Vec<_>>()
.join(",")
)
.fetch_optional(&mut *tx)
.execute(&mut *tx)
.await?;
// set_session_context(
// username TEXT,
// groups TEXT,
// pgroups TEXT,
// folders_read TEXT,
// folders_write TEXT
// )
// sqlx::query!(
// "SELECT set_config('session.user', $1, true)",
// authed.username()
// )
// .fetch_optional(&mut *tx)
// .await?;
// sqlx::query!(
// "SELECT set_config('session.groups', $1, true)",
// &authed.groups().join(",")
// )
// .fetch_optional(&mut *tx)
// .await?;
// sqlx::query!(
// "SELECT set_config('session.pgroups', $1, true)",
// &authed
// .groups()
// .iter()
// .map(|x| format!("g/{}", x))
// .collect::<Vec<_>>()
// .join(",")
// )
// .fetch_optional(&mut *tx)
// .await?;
// sqlx::query!(
// "SELECT set_config('session.folders_read', $1, true)",
// folders_read
// .iter()
// .map(|x| x.0.clone())
// .collect::<Vec<_>>()
// .join(",")
// )
// .fetch_optional(&mut *tx)
// .await?;
// sqlx::query!(
// "SELECT set_config('session.folders_write', $1, true)",
// folders_write
// .iter()
// .map(|x| x.0.clone())
// .collect::<Vec<_>>()
// .join(",")
// )
// .fetch_optional(&mut *tx)
// .await?;
Ok(tx)
}
}
+5 -2
View File
@@ -60,8 +60,11 @@ pub fn initialize_tracing(
"RUST_LOG",
&format!("windmill={}", rust_log_env.as_ref().unwrap()),
)
}
let default_env_filter = if rust_log_env.is_ok_and(|x| x == "debug") {
} else if rust_log_env.as_ref().is_ok_and(|x| x == "sqlxdebug") {
std::env::set_var("RUST_LOG", "windmill=debug,sqlx=debug");
};
let default_env_filter = if rust_log_env.is_ok_and(|x| x == "debug" || x == "sqlxdebug") {
LevelFilter::DEBUG
} else {
LevelFilter::INFO