Compare commits

...

2 Commits

Author SHA1 Message Date
563c274d70 lint-fix
Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
2024-12-18 12:48:08 +01:00
a9fee67b44 blueprints: gate apply and discovery behing pg lock
Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
2024-12-18 12:36:41 +01:00

View File

@ -5,8 +5,9 @@ from hashlib import sha512
from pathlib import Path from pathlib import Path
from sys import platform from sys import platform
import pglock
from dacite.core import from_dict from dacite.core import from_dict
from django.db import DatabaseError, InternalError, ProgrammingError from django.db import DatabaseError, InternalError, ProgrammingError, connection
from django.utils.text import slugify from django.utils.text import slugify
from django.utils.timezone import now from django.utils.timezone import now
from django.utils.translation import gettext_lazy as _ from django.utils.translation import gettext_lazy as _
@ -152,15 +153,27 @@ def blueprints_find() -> list[BlueprintFile]:
@prefill_task @prefill_task
def blueprints_discovery(self: SystemTask, path: str | None = None): def blueprints_discovery(self: SystemTask, path: str | None = None):
"""Find blueprints and check if they need to be created in the database""" """Find blueprints and check if they need to be created in the database"""
count = 0 with pglock.advisory(
for blueprint in blueprints_find(): lock_id=f"goauthentik.io/{connection.schema_name}/blueprints/discovery",
if path and blueprint.path != path: timeout=0,
continue side_effect=pglock.Return,
check_blueprint_v1_file(blueprint) ) as lock_acquired:
count += 1 if not lock_acquired:
self.set_status( LOGGER.debug("Not running blueprint discovery, lock was not acquired")
TaskStatus.SUCCESSFUL, _("Successfully imported {count} files.".format(count=count)) self.set_status(
) TaskStatus.SUCCESSFUL,
_("Blueprint discovery lock could not be acquired. Skipping discovery."),
)
return
count = 0
for blueprint in blueprints_find():
if path and blueprint.path != path:
continue
check_blueprint_v1_file(blueprint)
count += 1
self.set_status(
TaskStatus.SUCCESSFUL, _("Successfully imported {count} files.".format(count=count))
)
def check_blueprint_v1_file(blueprint: BlueprintFile): def check_blueprint_v1_file(blueprint: BlueprintFile):
@ -197,48 +210,60 @@ def check_blueprint_v1_file(blueprint: BlueprintFile):
def apply_blueprint(self: SystemTask, instance_pk: str): def apply_blueprint(self: SystemTask, instance_pk: str):
"""Apply single blueprint""" """Apply single blueprint"""
self.save_on_success = False self.save_on_success = False
instance: BlueprintInstance | None = None with pglock.advisory(
try: lock_id=f"goauthentik.io/{connection.schema_name}/blueprints/apply/{instance_pk}",
instance: BlueprintInstance = BlueprintInstance.objects.filter(pk=instance_pk).first() timeout=0,
if not instance or not instance.enabled: side_effect=pglock.Return,
) as lock_acquired:
if not lock_acquired:
LOGGER.debug("Not running blueprint discovery, lock was not acquired")
self.set_status(
TaskStatus.SUCCESSFUL,
_("Blueprint apply lock could not be acquired. Skipping apply."),
)
return return
self.set_uid(slugify(instance.name)) instance: BlueprintInstance | None = None
blueprint_content = instance.retrieve() try:
file_hash = sha512(blueprint_content.encode()).hexdigest() instance: BlueprintInstance = BlueprintInstance.objects.filter(pk=instance_pk).first()
importer = Importer.from_string(blueprint_content, instance.context) if not instance or not instance.enabled:
if importer.blueprint.metadata: return
instance.metadata = asdict(importer.blueprint.metadata) self.set_uid(slugify(instance.name))
valid, logs = importer.validate() blueprint_content = instance.retrieve()
if not valid: file_hash = sha512(blueprint_content.encode()).hexdigest()
instance.status = BlueprintInstanceStatus.ERROR importer = Importer.from_string(blueprint_content, instance.context)
instance.save() if importer.blueprint.metadata:
self.set_status(TaskStatus.ERROR, *logs) instance.metadata = asdict(importer.blueprint.metadata)
return valid, logs = importer.validate()
with capture_logs() as logs: if not valid:
applied = importer.apply()
if not applied:
instance.status = BlueprintInstanceStatus.ERROR instance.status = BlueprintInstanceStatus.ERROR
instance.save() instance.save()
self.set_status(TaskStatus.ERROR, *logs) self.set_status(TaskStatus.ERROR, *logs)
return return
instance.status = BlueprintInstanceStatus.SUCCESSFUL with capture_logs() as logs:
instance.last_applied_hash = file_hash applied = importer.apply()
instance.last_applied = now() if not applied:
self.set_status(TaskStatus.SUCCESSFUL) instance.status = BlueprintInstanceStatus.ERROR
except ( instance.save()
OSError, self.set_status(TaskStatus.ERROR, *logs)
DatabaseError, return
ProgrammingError, instance.status = BlueprintInstanceStatus.SUCCESSFUL
InternalError, instance.last_applied_hash = file_hash
BlueprintRetrievalFailed, instance.last_applied = now()
EntryInvalidError, self.set_status(TaskStatus.SUCCESSFUL)
) as exc: except (
if instance: OSError,
instance.status = BlueprintInstanceStatus.ERROR DatabaseError,
self.set_error(exc) ProgrammingError,
finally: InternalError,
if instance: BlueprintRetrievalFailed,
instance.save() EntryInvalidError,
) as exc:
if instance:
instance.status = BlueprintInstanceStatus.ERROR
self.set_error(exc)
finally:
if instance:
instance.save()
@CELERY_APP.task() @CELERY_APP.task()