feat: clean up auth method and include hostname in scheduler ctx

This commit is contained in:
Aravinth Manivannan 2023-10-20 01:41:46 +05:30
parent c8ecd29e94
commit 3e5dca9069
Signed by: realaravinth
GPG Key ID: F8F50389936984FF
1 changed files with 51 additions and 78 deletions

View File

@ -78,14 +78,7 @@ impl Data {
} }
/// Authenticate an mCaptcha instance and return its URL /// Authenticate an mCaptcha instance and return its URL
pub async fn mcaptcha_authenticate_and_get_url( pub async fn mcaptcha_authenticate(&self, secret: &str) -> ServiceResult<()> {
&self,
secret: &str,
) -> ServiceResult<Url> {
struct U {
url: String,
}
let res = sqlx::query!( let res = sqlx::query!(
"SELECT EXISTS ( "SELECT EXISTS (
SELECT SELECT
@ -99,27 +92,11 @@ impl Data {
.fetch_one(&self.db) .fetch_one(&self.db)
.await?; .await?;
if !match res.exists { if !matches!(res.exists, Some(true)) {
Some(true) => true,
_ => false,
} {
return Err(ServiceError::WrongPassword); return Err(ServiceError::WrongPassword);
} }
let url = sqlx::query_as!( Ok(())
U,
"SELECT
url
FROM
survey_mcaptcha_hostname
WHERE
secret = $1; ",
secret
)
.fetch_one(&self.db)
.await?;
Ok(Url::parse(&url.url).unwrap())
} }
/// Delete mCaptcha instance from database /// Delete mCaptcha instance from database
@ -264,7 +241,6 @@ impl Data {
pub async fn mcaptcha_get_checkpoint( pub async fn mcaptcha_get_checkpoint(
&self, &self,
campaign_id: &Uuid, campaign_id: &Uuid,
secret: &str,
) -> ServiceResult<usize> { ) -> ServiceResult<usize> {
let campaign_str = campaign_id.to_string(); let campaign_str = campaign_id.to_string();
@ -279,13 +255,8 @@ impl Data {
FROM FROM
survey_mcaptcha_campaign survey_mcaptcha_campaign
WHERE WHERE
campaign_id = $1 campaign_id = $1;",
AND
url_id = (
SELECT ID FROM survey_mcaptcha_hostname WHERE secret = $2
);",
&campaign_str, &campaign_str,
secret
) )
.fetch_one(&self.db) .fetch_one(&self.db)
.await?; .await?;
@ -297,7 +268,6 @@ impl Data {
pub async fn mcaptcha_set_checkpoint( pub async fn mcaptcha_set_checkpoint(
&self, &self,
campaign_id: &Uuid, campaign_id: &Uuid,
secret: &str,
checkpoint: usize, checkpoint: usize,
) -> ServiceResult<()> { ) -> ServiceResult<()> {
let campaign_str = campaign_id.to_string(); let campaign_str = campaign_id.to_string();
@ -307,15 +277,9 @@ impl Data {
SET SET
synced_till = $1 synced_till = $1
WHERE WHERE
campaign_id = $2 campaign_id = $2; ",
AND
url_id = (
SELECT ID FROM survey_mcaptcha_hostname WHERE secret = $3
)
",
checkpoint as i32, checkpoint as i32,
&campaign_str, &campaign_str,
secret
) )
.execute(&self.db) .execute(&self.db)
.await?; .await?;
@ -327,7 +291,6 @@ impl Data {
pub async fn mcaptcha_insert_analytics( pub async fn mcaptcha_insert_analytics(
&self, &self,
campaign_id: &Uuid, campaign_id: &Uuid,
secret: &str,
r: &PerformanceAnalytics, r: &PerformanceAnalytics,
) -> ServiceResult<()> { ) -> ServiceResult<()> {
let campaign_str = campaign_id.to_string(); let campaign_str = campaign_id.to_string();
@ -343,14 +306,9 @@ impl Data {
survey_mcaptcha_campaign survey_mcaptcha_campaign
WHERE WHERE
campaign_id = $1 campaign_id = $1
AND ), $2, $3, $4
url_id = (
SELECT ID FROM survey_mcaptcha_hostname WHERE secret = $2
)
), $3, $4, $5
);", );",
&campaign_str, &campaign_str,
secret,
r.time as i32, r.time as i32,
r.difficulty_factor as i32, r.difficulty_factor as i32,
&r.worker_type, &r.worker_type,
@ -413,18 +371,23 @@ impl Data {
Ok(res) Ok(res)
} }
pub async fn get_next_job_to_run(&self) -> ServiceResult<SchedulerJob> { pub async fn get_next_job_to_run(&self) -> ServiceResult<Option<SchedulerJob>> {
let res = sqlx::query_as!( let res = match sqlx::query_as!(
InnerSchedulerJob, InnerSchedulerJob,
"SELECT "SELECT
survey_mcaptcha_campaign.campaign_id, survey_mcaptcha_campaign.campaign_id,
survey_mcaptcha_upload_jobs.public_id survey_mcaptcha_upload_jobs.public_id,
survey_mcaptcha_hostname.url
FROM FROM
survey_mcaptcha_campaign survey_mcaptcha_campaign
INNER JOIN INNER JOIN
survey_mcaptcha_upload_jobs survey_mcaptcha_upload_jobs
ON ON
survey_mcaptcha_upload_jobs.campaign_id = survey_mcaptcha_campaign.ID survey_mcaptcha_upload_jobs.campaign_id = survey_mcaptcha_campaign.ID
INNER JOIN
survey_mcaptcha_hostname
ON
survey_mcaptcha_hostname.ID = survey_mcaptcha_campaign.url_id
WHERE WHERE
survey_mcaptcha_upload_jobs.job_state = ( survey_mcaptcha_upload_jobs.job_state = (
SELECT ID FROM survey_mcaptcha_upload_job_states WHERE name = $1 SELECT ID FROM survey_mcaptcha_upload_job_states WHERE name = $1
@ -437,8 +400,13 @@ impl Data {
&JOB_STATE_CREATE.name &JOB_STATE_CREATE.name
) )
.fetch_one(&self.db) .fetch_one(&self.db)
.await?; .await
Ok(res.into()) {
Ok(res) => Ok(Some(res.into())),
Err(sqlx::Error::RowNotFound) => Ok(None),
Err(e) => Err(e),
}?;
Ok(res)
} }
pub async fn add_job(&self, campaign_id: &Uuid) -> ServiceResult<Uuid> { pub async fn add_job(&self, campaign_id: &Uuid) -> ServiceResult<Uuid> {
@ -547,7 +515,7 @@ impl Data {
Err(e) => Err(e), Err(e) => Err(e),
}?; }?;
Ok(res.into()) Ok(res)
} }
pub async fn get_all_jobs_of_state( pub async fn get_all_jobs_of_state(
@ -588,7 +556,7 @@ impl Data {
Ok(res) Ok(res)
} }
pub async fn mark_job_scheduled(&self, public_id: &Uuid) -> ServiceResult<()> { pub async fn mark_job_scheduled(&self, job: &SchedulerJob) -> ServiceResult<()> {
let now = now_unix_time_stamp(); let now = now_unix_time_stamp();
sqlx::query!( sqlx::query!(
" "
@ -600,7 +568,7 @@ impl Data {
WHERE public_id = $3;", WHERE public_id = $3;",
&JOB_STATE_RUNNING.name, &JOB_STATE_RUNNING.name,
now, now,
&public_id.to_string(), &job.public_job_id.to_string(),
) )
.execute(&self.db) .execute(&self.db)
.await .await
@ -609,7 +577,7 @@ impl Data {
Ok(()) Ok(())
} }
pub async fn mark_job_finished(&self, public_id: &Uuid) -> ServiceResult<()> { pub async fn mark_job_finished(&self, job: &SchedulerJob) -> ServiceResult<()> {
let now = now_unix_time_stamp(); let now = now_unix_time_stamp();
sqlx::query!( sqlx::query!(
" "
@ -621,7 +589,7 @@ impl Data {
WHERE public_id = $3;", WHERE public_id = $3;",
&JOB_STATE_FINISH.name, &JOB_STATE_FINISH.name,
now, now,
&public_id.to_string(), &job.public_job_id.to_string(),
) )
.execute(&self.db) .execute(&self.db)
.await .await
@ -635,18 +603,21 @@ impl Data {
pub struct SchedulerJob { pub struct SchedulerJob {
pub campaign_id: Uuid, pub campaign_id: Uuid,
pub public_job_id: Uuid, pub public_job_id: Uuid,
pub url: Url,
} }
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
struct InnerSchedulerJob { struct InnerSchedulerJob {
campaign_id: String, campaign_id: String,
public_id: String, public_id: String,
url: String,
} }
impl From<InnerSchedulerJob> for SchedulerJob { impl From<InnerSchedulerJob> for SchedulerJob {
fn from(j: InnerSchedulerJob) -> Self { fn from(j: InnerSchedulerJob) -> Self {
SchedulerJob { SchedulerJob {
campaign_id: Uuid::parse_str(&j.campaign_id).unwrap(), campaign_id: Uuid::parse_str(&j.campaign_id).unwrap(),
public_job_id: Uuid::parse_str(&j.public_id).unwrap(), public_job_id: Uuid::parse_str(&j.public_id).unwrap(),
url: Url::parse(&j.url).unwrap(),
} }
} }
} }
@ -721,11 +692,10 @@ mod tests {
assert_ne!(secret2, secret); assert_ne!(secret2, secret);
let secret = secret2; let secret = secret2;
assert!(data.mcaptcha_authenticate(&secret).await.is_ok());
assert_eq!( assert_eq!(
data.mcaptcha_authenticate_and_get_url(&secret) data.mcaptcha_authenticate("foo").await.err(),
.await Some(ServiceError::WrongPassword)
.unwrap(),
url
); );
let uuid = Uuid::new_v4(); let uuid = Uuid::new_v4();
@ -752,17 +722,9 @@ mod tests {
.await .await
.unwrap()); .unwrap());
assert_eq!( assert_eq!(data.mcaptcha_get_checkpoint(&uuid).await.unwrap(), 0);
data.mcaptcha_get_checkpoint(&uuid, &secret).await.unwrap(), data.mcaptcha_set_checkpoint(&uuid, 1).await.unwrap();
0 assert_eq!(data.mcaptcha_get_checkpoint(&uuid).await.unwrap(), 1);
);
data.mcaptcha_set_checkpoint(&uuid, &secret, 1)
.await
.unwrap();
assert_eq!(
data.mcaptcha_get_checkpoint(&uuid, &secret).await.unwrap(),
1
);
let analytics = PerformanceAnalytics { let analytics = PerformanceAnalytics {
id: 1, id: 1,
@ -770,7 +732,7 @@ mod tests {
difficulty_factor: 1, difficulty_factor: 1,
worker_type: "foo".to_string(), worker_type: "foo".to_string(),
}; };
data.mcaptcha_insert_analytics(&uuid, &secret, &analytics) data.mcaptcha_insert_analytics(&uuid, &analytics)
.await .await
.unwrap(); .unwrap();
@ -813,9 +775,15 @@ mod tests {
let job2_public_id = data.add_job(&uuid).await.unwrap(); let job2_public_id = data.add_job(&uuid).await.unwrap();
let job2 = data.get_job(&job2_public_id).await.unwrap().unwrap(); let job2 = data.get_job(&job2_public_id).await.unwrap().unwrap();
assert_eq!(job2, job); assert_eq!(job2, job);
let scheduler_job = data.get_next_job_to_run().await.unwrap().unwrap();
assert_eq!(scheduler_job.url, url);
assert_eq!( assert_eq!(
data.get_next_job_to_run().await.unwrap().public_job_id, data.get_next_job_to_run()
.await
.unwrap()
.unwrap()
.public_job_id,
job.public_job_id job.public_job_id
); );
@ -827,7 +795,8 @@ mod tests {
vec![job.clone()] vec![job.clone()]
); );
data.mark_job_scheduled(&job.public_job_id).await.unwrap(); data.mark_job_scheduled(&scheduler_job).await.unwrap();
assert!(data.get_next_job_to_run().await.unwrap().is_none(),);
let job = data.get_job(&job.public_job_id).await.unwrap().unwrap(); let job = data.get_job(&job.public_job_id).await.unwrap().unwrap();
assert!(job.scheduled_at.is_some()); assert!(job.scheduled_at.is_some());
assert_eq!( assert_eq!(
@ -837,7 +806,7 @@ mod tests {
vec![job.clone()] vec![job.clone()]
); );
data.mark_job_finished(&job.public_job_id).await.unwrap(); data.mark_job_finished(&scheduler_job).await.unwrap();
let job = data.get_job(&job.public_job_id).await.unwrap().unwrap(); let job = data.get_job(&job.public_job_id).await.unwrap().unwrap();
assert!(job.finished_at.is_some()); assert!(job.finished_at.is_some());
assert_eq!( assert_eq!(
@ -849,7 +818,11 @@ mod tests {
let job2 = data.get_job(&job2_public_id).await.unwrap().unwrap(); let job2 = data.get_job(&job2_public_id).await.unwrap().unwrap();
assert_ne!(job2.public_job_id, job.public_job_id); assert_ne!(job2.public_job_id, job.public_job_id);
assert_eq!( assert_eq!(
data.get_next_job_to_run().await.unwrap().public_job_id, data.get_next_job_to_run()
.await
.unwrap()
.unwrap()
.public_job_id,
job2.public_job_id job2.public_job_id
); );
assert_eq!(public_id, job2.campaign_public_id); assert_eq!(public_id, job2.campaign_public_id);