feat: add WHITELIST_WORKSPACES and BLACKLIST_WORKSPACES

This commit is contained in:
Ruben Fiszel
2023-02-02 07:57:03 +01:00
parent 68500b12b2
commit 99568eaa47
3 changed files with 40 additions and 7 deletions
+2
View File
@@ -312,6 +312,8 @@ upcoming CLI tool.
| PIP_LOCAL_DEPENDENCIES | None | Specify dependencies that are installed locally and do not need to be solved nor installed again |
| ADDITIONAL_PYTHON_PATHS | None | Specify python paths (separated by a :) to be appended to the PYTHONPATH of the python jobs. To be used with PIP_LOCAL_DEPENDENCIES to use python codebases within Windmill | Worker |
| INCLUDE_HEADERS | None | Whitelist of headers that are passed to jobs as args (separated by a comma) | Server |
| WHITELIST_WORKSPACES | None | Whitelist of workspaces this worker takes job from | Worker |
| BLACKLIST_WORKSPACES | None | Blacklist of workspaces this worker takes job from | Worker |
+31 -6
View File
@@ -82,7 +82,32 @@ pub async fn cancel_job<'c>(
Ok((tx, job_option))
}
pub async fn pull(db: &Pool<Postgres>) -> windmill_common::error::Result<Option<QueuedJob>> {
pub async fn pull(
db: &Pool<Postgres>,
whitelist_workspaces: Option<Vec<String>>,
blacklist_workspaces: Option<Vec<String>>,
) -> windmill_common::error::Result<Option<QueuedJob>> {
let mut workspaces_filter = String::new();
if let Some(whitelist) = whitelist_workspaces {
workspaces_filter.push_str(&format!(
" AND workspace_id IN ({})",
whitelist
.into_iter()
.map(|x| format!("'{x}'"))
.collect::<Vec<String>>()
.join(",")
));
}
if let Some(blacklist) = blacklist_workspaces {
workspaces_filter.push_str(&format!(
" AND workspace_id NOT IN ({})",
blacklist
.into_iter()
.map(|x| format!("'{x}'"))
.collect::<Vec<String>>()
.join(",")
));
}
/* Jobs can be started if they:
* - haven't been started before,
* running = false
@@ -90,7 +115,7 @@ pub async fn pull(db: &Pool<Postgres>) -> windmill_common::error::Result<Option<
* suspend_until is non-null
* and suspend = 0 when the resume messages are received
* or suspend_until <= now() if it has timed out */
let job: Option<QueuedJob> = sqlx::query_as::<_, QueuedJob>(
let job: Option<QueuedJob> = sqlx::query_as::<_, QueuedJob>(&format!(
"UPDATE queue
SET running = true
, started_at = coalesce(started_at, now())
@@ -99,17 +124,17 @@ pub async fn pull(db: &Pool<Postgres>) -> windmill_common::error::Result<Option<
WHERE id = (
SELECT id
FROM queue
WHERE ( running = false
WHERE ((running = false
AND scheduled_for <= now())
OR (suspend_until IS NOT NULL
AND ( suspend <= 0
OR suspend_until <= now()))
OR suspend_until <= now()))) {workspaces_filter}
ORDER BY scheduled_for
FOR UPDATE SKIP LOCKED
LIMIT 1
)
RETURNING *",
)
RETURNING *"
))
.fetch_optional(db)
.await?;
+7 -1
View File
@@ -430,6 +430,12 @@ pub async fn run_worker(
let pip_local_dependencies = std::env::var("PIP_LOCAL_DEPENDENCIES")
.ok()
.map(|x| x.split(',').map(|x| x.to_string()).collect());
let whitelist_workspaces = std::env::var("WHITELIST_WORKSPACES")
.ok()
.map(|x| x.split(',').map(|x| x.to_string()).collect());
let blacklist_workspaces = std::env::var("BLACKLIST_WORKSPACES")
.ok()
.map(|x| x.split(',').map(|x| x.to_string()).collect());
let pip_local_dependencies = if pip_local_dependencies == Some(vec!["".to_string()]) {
None
@@ -526,7 +532,7 @@ pub async fn run_worker(
.await
.map_err(|_| Error::InternalErr("Impossible to fetch same_worker job".to_string())))
},
(job, timer) = {let timer = worker_pull_duration.start_timer(); pull(&db).map(|x| (x, timer)) } => {
(job, timer) = {let timer = worker_pull_duration.start_timer(); pull(&db, whitelist_workspaces.clone(), blacklist_workspaces.clone()).map(|x| (x, timer)) } => {
drop(timer);
(false, job)
},