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