Add state management to job (#1113)

* Add state management to job

* Add migration

Co-authored-by: Adrià Casajús <adria.casajus@proton.ch>
This commit is contained in:
Adrià Casajús 2022-06-28 09:22:48 +02:00 committed by GitHub
parent 686f4f3f68
commit 5fa41d6ccf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 172 additions and 113 deletions

View File

@ -262,6 +262,13 @@ class VerpType(EnumE):
transactional = 2
class JobState(EnumE):
ready = 0
taken = 1
done = 2
error = 3
class Hibp(Base, ModelMixin):
__tablename__ = "hibp"
name = sa.Column(sa.String(), nullable=False, unique=True, index=True)
@ -2370,6 +2377,14 @@ class Job(Base, ModelMixin):
# whether the job has been taken by the job runner
taken = sa.Column(sa.Boolean, default=False, nullable=False)
run_at = sa.Column(ArrowType)
state = sa.Column(
sa.Integer,
nullable=False,
server_default=str(JobState.ready.value),
default=JobState.ready.value,
)
attempts = sa.Column(sa.Integer, nullable=False, server_default="0", default=0)
taken_at = sa.Column(ArrowType, nullable=True)
def __repr__(self):
return f"<Job {self.id} {self.name} {self.payload}>"

View File

@ -15,7 +15,7 @@ from app.email_utils import (
from app.import_utils import handle_batch_import
from app.jobs.export_user_data_job import ExportUserDataJob
from app.log import LOG
from app.models import User, Job, BatchImport, Mailbox, CustomDomain
from app.models import User, Job, BatchImport, Mailbox, CustomDomain, JobState
from server import create_light_app
@ -106,6 +106,120 @@ def welcome_proton(user):
)
def process_job(job: Job):
if job.name == config.JOB_ONBOARDING_1:
user_id = job.payload.get("user_id")
user = User.get(user_id)
# user might delete their account in the meantime
# or disable the notification
if user and user.notification and user.activated:
LOG.d("send onboarding send-from-alias email to user %s", user)
onboarding_send_from_alias(user)
elif job.name == config.JOB_ONBOARDING_2:
user_id = job.payload.get("user_id")
user = User.get(user_id)
# user might delete their account in the meantime
# or disable the notification
if user and user.notification and user.activated:
LOG.d("send onboarding mailbox email to user %s", user)
onboarding_mailbox(user)
elif job.name == config.JOB_ONBOARDING_4:
user_id = job.payload.get("user_id")
user = User.get(user_id)
# user might delete their account in the meantime
# or disable the notification
if user and user.notification and user.activated:
LOG.d("send onboarding pgp email to user %s", user)
onboarding_pgp(user)
elif job.name == config.JOB_BATCH_IMPORT:
batch_import_id = job.payload.get("batch_import_id")
batch_import = BatchImport.get(batch_import_id)
handle_batch_import(batch_import)
elif job.name == config.JOB_DELETE_ACCOUNT:
user_id = job.payload.get("user_id")
user = User.get(user_id)
if not user:
LOG.i("No user found for %s", user_id)
return
user_email = user.email
LOG.w("Delete user %s", user)
User.delete(user.id)
Session.commit()
send_email(
user_email,
"Your SimpleLogin account has been deleted",
render("transactional/account-delete.txt"),
render("transactional/account-delete.html"),
retries=3,
)
elif job.name == config.JOB_DELETE_MAILBOX:
mailbox_id = job.payload.get("mailbox_id")
mailbox = Mailbox.get(mailbox_id)
if not mailbox:
return
mailbox_email = mailbox.email
user = mailbox.user
Mailbox.delete(mailbox_id)
Session.commit()
LOG.d("Mailbox %s %s deleted", mailbox_id, mailbox_email)
send_email(
user.email,
f"Your mailbox {mailbox_email} has been deleted",
f"""Mailbox {mailbox_email} along with its aliases are deleted successfully.
Regards,
SimpleLogin team.
""",
retries=3,
)
elif job.name == config.JOB_DELETE_DOMAIN:
custom_domain_id = job.payload.get("custom_domain_id")
custom_domain = CustomDomain.get(custom_domain_id)
if not custom_domain:
return
domain_name = custom_domain.domain
user = custom_domain.user
CustomDomain.delete(custom_domain.id)
Session.commit()
LOG.d("Domain %s deleted", domain_name)
send_email(
user.email,
f"Your domain {domain_name} has been deleted",
f"""Domain {domain_name} along with its aliases are deleted successfully.
Regards,
SimpleLogin team.
""",
retries=3,
)
elif job.name == config.JOB_SEND_USER_REPORT:
export_job = ExportUserDataJob.create_from_job(job)
if export_job:
export_job.run()
elif job.name == config.JOB_SEND_PROTON_WELCOME_1:
user_id = job.payload.get("user_id")
user = User.get(user_id)
if user and user.activated:
LOG.d("send proton welcome email to user %s", user)
welcome_proton(user)
else:
LOG.e("Unknown job name %s", job.name)
if __name__ == "__main__":
while True:
# wrap in an app context to benefit from app setup like database cleanup, sentry integration, etc
@ -114,6 +228,8 @@ if __name__ == "__main__":
min_dt = arrow.now().shift(hours=-1)
max_dt = arrow.now().shift(hours=1)
# TODO: Change This condition after deploying this MR
# to Job.state == ready or (Job.state == taken and job.taken_at < arrow.now.shift(minutes=-10))
for job in Job.filter(
Job.taken.is_(False), Job.run_at > min_dt, Job.run_at <= max_dt
).all():
@ -121,118 +237,13 @@ if __name__ == "__main__":
# mark the job as taken, whether it will be executed successfully or not
job.taken = True
job.taken_at = arrow.now()
job.state = JobState.taken.value
job.attempts += 1
Session.commit()
process_job(job)
job.state = JobState.done.value
Session.commit()
if job.name == config.JOB_ONBOARDING_1:
user_id = job.payload.get("user_id")
user = User.get(user_id)
# user might delete their account in the meantime
# or disable the notification
if user and user.notification and user.activated:
LOG.d("send onboarding send-from-alias email to user %s", user)
onboarding_send_from_alias(user)
elif job.name == config.JOB_ONBOARDING_2:
user_id = job.payload.get("user_id")
user = User.get(user_id)
# user might delete their account in the meantime
# or disable the notification
if user and user.notification and user.activated:
LOG.d("send onboarding mailbox email to user %s", user)
onboarding_mailbox(user)
elif job.name == config.JOB_ONBOARDING_4:
user_id = job.payload.get("user_id")
user = User.get(user_id)
# user might delete their account in the meantime
# or disable the notification
if user and user.notification and user.activated:
LOG.d("send onboarding pgp email to user %s", user)
onboarding_pgp(user)
elif job.name == config.JOB_BATCH_IMPORT:
batch_import_id = job.payload.get("batch_import_id")
batch_import = BatchImport.get(batch_import_id)
handle_batch_import(batch_import)
elif job.name == config.JOB_DELETE_ACCOUNT:
user_id = job.payload.get("user_id")
user = User.get(user_id)
if not user:
LOG.i("No user found for %s", user_id)
continue
user_email = user.email
LOG.w("Delete user %s", user)
User.delete(user.id)
Session.commit()
send_email(
user_email,
"Your SimpleLogin account has been deleted",
render("transactional/account-delete.txt"),
render("transactional/account-delete.html"),
retries=3,
)
elif job.name == config.JOB_DELETE_MAILBOX:
mailbox_id = job.payload.get("mailbox_id")
mailbox = Mailbox.get(mailbox_id)
if not mailbox:
continue
mailbox_email = mailbox.email
user = mailbox.user
Mailbox.delete(mailbox_id)
Session.commit()
LOG.d("Mailbox %s %s deleted", mailbox_id, mailbox_email)
send_email(
user.email,
f"Your mailbox {mailbox_email} has been deleted",
f"""Mailbox {mailbox_email} along with its aliases are deleted successfully.
Regards,
SimpleLogin team.
""",
retries=3,
)
elif job.name == config.JOB_DELETE_DOMAIN:
custom_domain_id = job.payload.get("custom_domain_id")
custom_domain = CustomDomain.get(custom_domain_id)
if not custom_domain:
continue
domain_name = custom_domain.domain
user = custom_domain.user
CustomDomain.delete(custom_domain.id)
Session.commit()
LOG.d("Domain %s deleted", domain_name)
send_email(
user.email,
f"Your domain {domain_name} has been deleted",
f"""Domain {domain_name} along with its aliases are deleted successfully.
Regards,
SimpleLogin team.
""",
retries=3,
)
elif job.name == config.JOB_SEND_USER_REPORT:
export_job = ExportUserDataJob.create_from_job(job)
if export_job:
export_job.run()
elif job.name == config.JOB_SEND_PROTON_WELCOME_1:
user_id = job.payload.get("user_id")
user = User.get(user_id)
if user and user.activated:
LOG.d("send proton welcome email to user %s", user)
welcome_proton(user)
else:
LOG.e("Unknown job name %s", job.name)
time.sleep(10)

View File

@ -0,0 +1,33 @@
"""Add state to job
Revision ID: bfebc2d5c719
Revises: d1fb679f7eec
Create Date: 2022-06-27 14:56:58.797121
"""
import sqlalchemy_utils
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = 'bfebc2d5c719'
down_revision = 'd1fb679f7eec'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('job', sa.Column('attempts', sa.Integer(), server_default='0', nullable=False))
op.add_column('job', sa.Column('state', sa.Integer(), server_default='0', nullable=False))
op.add_column('job', sa.Column('taken_at', sqlalchemy_utils.types.arrow.ArrowType(), nullable=True))
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column('job', 'taken_at')
op.drop_column('job', 'state')
op.drop_column('job', 'attempts')
# ### end Alembic commands ###